torchmil is a PyTorch-based library for deep Multiple Instance Learning (MIL). It provides a simple, flexible, and extensible framework for working with MIL models and data.
It includes:
- A collection of popular MIL models.
- Different PyTorch modules frequently used in MIL models.
- Handy tools to deal with MIL data.
- A collection of popular MIL datasets.
pip install torchmil
You can load a MIL dataset and train a MIL model in just a few lines of code:
from torchmil.datasets import Camelyon16MIL
from torchmil.models import ABMIL
from torchmil.utils import Trainer
from torchmil.data import collate_fn
from torch.utils.data import DataLoader
# Load the Camelyon16 dataset
dataset = Camelyon16MIL(root='data', features='UNI')
dataloader = DataLoader(dataset, batch_size=4, shuffle=True, collate_fn=collate_fn)
# Instantiate the ABMIL model and optimizer
model = ABMIL(in_shape=(2048,), criterion=torch.nn.BCEWithLogitsLoss()) # each model has its own criterion
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# Instantiate the Trainer
trainer = Trainer(model, optimizer, device='cuda')
# Train the model
trainer.train(dataloader, epochs=10)
# Save the model
torch.save(model.state_dict(), 'model.pth')
You can take a look at the examples to see how to use torchmil in practice. To see the full list of available models, datasets, and modules, check the API reference.
We welcome contributions to torchmil! There several ways you can contribute:
- Reporting bugs or issues you encounter while using the library, asking questions, or requesting new features: use the Github issues.
- Improving the documentation: if you find any part of the documentation unclear or incomplete, feel free to submit a pull request with improvements.
- If you have a new model, dataset, or utility that you think would be useful for the community, please consider submitting a pull request to add it to the library.
Take a look at CONTRIBUTING.md for more details on how to contribute.
If you find this library useful, please consider citing it:
@misc{torchmil,
author = {Castro-Mac{\'\i}as, Francisco M and S{\'a}ez-Maldonado, Francisco Javier and Morales Alvarez, Pablo and Molina, Rafael},
title = {torchmil: A PyTorch-based library for deep Multiple Instance Learning},
year = {2025},
howpublished = {\url{https://franblueee.github.io/torchmil/}}
}