Lightning API

The datamint.lightning module provides PyTorch Lightning integration for training machine learning models with Datamint datasets.

Lightning Components

DatamintDataModule

The DatamintDataModule wraps any DatamintBaseDataset and provides train_dataloader, val_dataloader, test_dataloader, and predict_dataloader for use with a Lightning Trainer.

import albumentations as A
from datamint.dataset import ImageDataset
from datamint.lightning import DatamintDataModule

train_tfm = A.Compose([A.RandomHorizontalFlip(), A.Normalize()])
eval_tfm  = A.Compose([A.Normalize()])

dataset = ImageDataset(project='my_project')
dm = DatamintDataModule(
    dataset,
    batch_size=8,
    split={'train': 0.8, 'val': 0.1, 'test': 0.1},
    split_seed=42,
    train_transform=train_tfm,
    eval_transform=eval_tfm,
)

trainer = lightning.Trainer(...)
trainer.fit(model, datamodule=dm)
trainer.test(datamodule=dm)

Constructor Parameters

DatamintDataModule — LightningDataModule wrapper for Datamint datasets.

Wraps any DatamintBaseDataset subclass and provides train_dataloader, val_dataloader, test_dataloader, and predict_dataloader for use with a Lightning Trainer.

Key Methods

DatamintDataModule — LightningDataModule wrapper for Datamint datasets.

Wraps any DatamintBaseDataset subclass and provides train_dataloader, val_dataloader, test_dataloader, and predict_dataloader for use with a Lightning Trainer.

Trainers

The trainer layer packages the usual Lightning workflow into a small number of task-focused entry points. A trainer can:

  • Build the dataset and datamodule for a Datamint project,

  • Choose task-specific default transforms, loss functions, and metrics,

  • Create the Lightning trainer, MLflow logger, and checkpoint callbacks,

  • Train and test the model, and

  • Optionally register the resulting model in MLflow.

Available Trainers

Specialized trainers for end-to-end Datamint workflows.

class datamint.lightning.trainers.BaseTrainer(dataset=None, project=None, *, dataset_kwargs=None, model=None, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, **kwargs)

Bases: ABC

Abstract base trainer encapsulating an end-to-end training workflow.

Subclasses provide task-specific defaults for model architecture, transforms, loss, and metrics by overriding the _build_* / _default_* hooks. Users typically only need to specify a project (or dataset) and optionally override a few settings.

Parameters:

  • dataset (DatamintBaseDataset | None) – A pre-built DatamintBaseDataset. Mutually exclusive with project.

  • project (str | Project | None) – Project name or Project object used to auto-build a dataset when dataset is None.

  • model (LightningModule | type[LightningModule] | None) – A user-provided LightningModule. When None the trainer builds a default one via _build_model().

  • loss_fn (Module | None) – Custom loss function forwarded to the default model. Ignored when model is provided (the user’s module owns its own loss).

  • batch_size (int) – Training batch size.

  • num_workers (int) – DataLoader workers.

  • train_transform (BaseCompose | None) – Albumentations transform for training. When None the trainer uses _train_transform().

  • eval_transform (BaseCompose | None) – Albumentations transform for val/test. When None the trainer uses _eval_transform().

  • split_as_of_timestamp (str | None) – Historical timestamp used to resolve project-scoped dataset splits during training. When omitted, the resolved project split datasets capture the current UTC timestamp and training lineage logs it via MLflow.

  • max_epochs (int) – Maximum number of training epochs.

  • early_stopping_patience (int | None) – Epochs without improvement before stopping. Set to None to disable early stopping.

  • mlflow_experiment_name (str | None) – MLflow experiment name. Auto-generated from the project name when None.

  • model_name (str | None) – Name for the model in the registry. Auto-generated when None.

  • auto_deploy_adapter (bool) – When True, auto-generate a DatamintModel adapter after training.

  • trainer_kwargs (dict[str, Any] | None) – Extra keyword arguments forwarded to lightning.Trainer.

  • dataset_kwargs (dict[str, Any] | None)

  • kwargs (Any)

property datamodule: DatamintDataModule
property dataset: DatamintBaseDataset
property experiment_name: str
fit()

Run the full training pipeline.

Return type:

dict[str, Any]

Returns:

Dictionary with keys 'trainer', 'model', 'test_results', and 'adapter' (when auto_deploy_adapter is enabled).

property model: LightningModule
test(register_model=True)

Run evaluation on the test split in a fresh run.

Parameters:

register_model (bool) – When True, run a zero-epoch fit first so the checkpoint callback saves the current model to MLflow and registers it after test metrics are logged.

Return type:

list[Mapping[str, float]]

class datamint.lightning.trainers.ClassificationTrainer(dataset=None, project=None, *, dataset_kwargs=None, model=None, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, **kwargs)

Bases: BaseTrainer

Abstract trainer for classification tasks.

Provides shared defaults:

  • Loss – CrossEntropyLoss.

  • Metrics – Multiclass Accuracy and macro F1 (torchmetrics).

  • Monitor – val/accuracy (maximise).

Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • dataset_kwargs (dict[str, Any] | None)

  • model (LightningModule | type[LightningModule] | None)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • kwargs (Any)

class datamint.lightning.trainers.DeepLabV3PlusTrainer(dataset=None, project=None, *, image_size=None, slice_axis=None, model=None, in_channels=3, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, dataset_kwargs=None, encoder_name='resnet34', decoder_atrous_rates=(12, 24, 36), **kwargs)

Bases: SemanticSegmentation2DTrainer

Convenience trainer pre-configured for DeepLab v3+.

Uses the ASPP-based DeepLab v3+ architecture from segmentation_models_pytorch. The decoder_atrous_rates parameter controls the dilation rates of the Atrous Spatial Pyramid Pooling module, which is DeepLab v3+’s core multi-scale context mechanism.

Example:

trainer = DeepLabV3PlusTrainer(
    project='BUS_Segmentation',
    encoder_name='resnet50',
)
results = trainer.fit()
Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • image_size (int | tuple[int, int] | None)

  • slice_axis (Literal['axial', 'sagittal', 'coronal'] | int | None)

  • model (LightningModule | type[LightningModule] | None)

  • in_channels (int)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • dataset_kwargs (dict[str, Any] | None)

  • encoder_name (str)

  • decoder_atrous_rates (tuple[int, int, int])

  • kwargs (Any)

class datamint.lightning.trainers.DetectionTrainer(dataset=None, project=None, *, dataset_kwargs=None, model=None, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, **kwargs)

Bases: BaseTrainer

Abstract trainer for object detection tasks.

Provides shared defaults for all detection models:

  • Dataset – ImageDataset with return_boxes=True

  • Collate – detection_collate_fn() (variable-length boxes)

  • Metrics – Mean Average Precision (torchmetrics)

  • Monitor – val/map (maximise)

Subclasses must implement _build_model(), _train_transform(), and _eval_transform().

Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • dataset_kwargs (dict[str, Any] | None)

  • model (LightningModule | type[LightningModule] | None)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • kwargs (Any)

class datamint.lightning.trainers.EfficientNetV2Trainer(*, architecture='efficientnetv2_s', image_size=384, **kwargs)

Bases: ImageClassificationTrainer

Trainer pre-configured for EfficientNetV2.

Default model: EfficientNetV2-S pretrained on ImageNet at 384×384.

Parameters:

  • architecture (str) – timm model name. Defaults to 'efficientnetv2_s'. Other valid choices: 'efficientnetv2_m', 'efficientnetv2_l', 'efficientnetv2_xl', 'efficientnetv2_rw_t'.

  • image_size (int | tuple[int, int]) – Target image size (H, W) or a single int for square images. Defaults to 384, the resolution EfficientNetV2-S was trained at.

Example:

trainer = EfficientNetV2Trainer(project='ChestXray')
results = trainer.fit()
Parameters:

kwargs (Any)

class datamint.lightning.trainers.ImageClassificationTrainer(*, architecture='resnet34', pretrained=True, image_size=None, **kwargs)

Bases: ClassificationTrainer

Trainer for image classification tasks.

Default model: ResNet-34 (via timm) pretrained on ImageNet.

Parameters:

  • architecture (str) – timm model name. Defaults to 'resnet34'.

  • pretrained (bool) – Use pretrained weights. Defaults to True.

  • image_size (int | tuple[int, int] | None) – Optional target image size (H, W) or a single int for square images. When omitted, the trainer keeps the original image size instead of forcing a resize.

Example:

trainer = ImageClassificationTrainer(project='ChestXray')
results = trainer.fit()
Parameters:

kwargs (Any)

class datamint.lightning.trainers.NNUNetTrainer(dataset=None, project=None, *, configuration='3d_fullres', fold=0, dataset_id=None, nnunet_work_dir=None, continue_training=False, channel_names=None, num_processes_preprocessing=None, max_epochs=1000, **kwargs)

Bases: BaseTrainer

Datamint trainer that runs nnUNet v2 as the training backend.

Exports the project data to nnUNet Task format, runs fingerprinting, planning, preprocessing, and training via _DatamintNNUNetTrainer, then imports predictions back as Datamint annotations.

Parameters:

  • project (str | Project | None) – Datamint project name or object.

  • configuration (str) – nnUNet configuration — '2d', '3d_fullres' (default), '3d_lowres', or '3d_cascade_fullres'.

  • fold (int | str) – Cross-validation fold index (0–4) or 'all'.

  • dataset_id (int | None) – Fixed nnUNet dataset ID (1–999). When None the ID is auto-assigned from the registry.

  • nnunet_work_dir (Path | str | None) – Root directory for all nnUNet I/O (nnUNet_raw, nnUNet_preprocessed, nnUNet_results are created beneath it). Defaults to ~/.cache/datamint/nnunet/.

  • continue_training (bool) – Resume from an existing checkpoint.

  • channel_names (dict[str, str] | None) – Modality index → name mapping passed to dataset.json, e.g. {'0': 'CT'}.

  • num_processes_preprocessing (int | None) – Workers for nnUNet’s preprocessing step. None lets nnUNet choose.

  • max_epochs (int) – Training epochs forwarded to nnUNetTrainer.

  • kwargs (Any)

fit()

Run the full nnUNet training pipeline.

Steps in order:

  1. Assign a stable nnUNet dataset ID for this project.

  2. Set the three nnUNet_* environment variables.

  3. Export all project resources to nnUNet Task format on disk.

  4. Run dataset fingerprinting and experiment planning.

  5. Run preprocessing for the configured nnUNet configuration.

  6. Start an MLflow run.

  7. Instantiate _DatamintNNUNetTrainer and call run_training().

  8. Run inference on the test split (_run_prediction).

  9. Import predictions as Datamint annotations (_import_predictions).

Return type:

dict

Returns:

{'bridge': bridge} — the trained bridge instance.

class datamint.lightning.trainers.SegmentationTrainer(dataset=None, project=None, *, dataset_kwargs=None, model=None, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, **kwargs)

Bases: BaseTrainer

Abstract trainer for segmentation tasks.

Provides shared defaults:

  • Loss – combined BCE + Dice (_BCEDiceLoss).

  • Metrics – Mean IoU and Generalised Dice Score (torchmetrics).

  • Monitor – val/iou (maximise).

Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • dataset_kwargs (dict[str, Any] | None)

  • model (LightningModule | type[LightningModule] | None)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • kwargs (Any)

class datamint.lightning.trainers.SemanticSegmentation2DTrainer(*, image_size=None, slice_axis=None, model=None, in_channels=3, trainer_kwargs=None, **kwargs)

Bases: SegmentationTrainer

Trainer for 2-D semantic segmentation.

Default model: UNet++ (segmentation_models_pytorch) with a resnet34 encoder pretrained on ImageNet.

When pointed at a project made of 3-D volumes, the trainer automatically converts it to a SlicedVolumeDataset and trains on 2-D slices instead.

Parameters:

  • slice_axis (Literal['axial', 'sagittal', 'coronal'] | int | None) – Slice axis override for 3-D volume projects. When omitted, the trainer tries to infer the most sensible anatomical plane and falls back to 'axial'.

  • image_size (int | tuple[int, int] | None) – Target image size (H, W) or a single int for square images. Forwarded to default transforms. When None a sensible default is chosen.

  • in_channels (int) – Number of input image channels. Defaults to 3.

  • to (All remaining keyword arguments are forwarded)

:param BaseTrainer.:

Example:

trainer = SemanticSegmentation2DTrainer(project='BUS_Segmentation')
results = trainer.fit()
Parameters:

  • model (LightningModule | type[LightningModule] | None)

  • trainer_kwargs (dict[str, Any] | None)

  • kwargs (Any)

slice_axis: Literal['axial', 'sagittal', 'coronal'] | int | None
class datamint.lightning.trainers.SemanticSegmentation3DTrainer(*, slice_axis='axial', encoder_name='resnet34', in_channels=3, image_size=None, **kwargs)

Bases: SegmentationTrainer

Trainer for 3-D semantic segmentation via per-slice 2-D training.

Builds a VolumeDataset, slices it along the chosen axis, and trains a 2-D segmentation model on individual slices.

Parameters:

  • slice_axis (str | int) – Slicing axis — 'axial', 'sagittal', 'coronal', or an integer axis index.

  • encoder_name (str) – SMP encoder backbone.

  • in_channels (int) – Number of input channels.

  • image_size (int | tuple[int, int] | None) – Optional target image size (H, W) or a single int for square images. When omitted, training keeps the original slice size.

Example:

trainer = SemanticSegmentation3DTrainer(
    project='CT_Liver',
    slice_axis='axial',
)
results = trainer.fit()
Parameters:

kwargs (Any)

class datamint.lightning.trainers.TransUNetTrainer(dataset=None, project=None, *, image_size=None, slice_axis=None, model=None, in_channels=3, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, dataset_kwargs=None, variant='R50-ViT-B_16', pretrained=True, **kwargs)

Bases: SemanticSegmentation2DTrainer

Convenience trainer pre-configured for TransUNet.

Uses the R50-ViT-B/16 hybrid encoder with a Cascaded UPsampler (CUP) decoder from Chen et al. (2021). The backbone is timm’s vit_base_r50_s16_224, which is a drop-in match for the architecture described in the paper.

Example:

trainer = TransUNetTrainer(
    project='BUS_Segmentation',
)
results = trainer.fit()
Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • image_size (int | tuple[int, int] | None)

  • slice_axis (Literal['axial', 'sagittal', 'coronal'] | int | None)

  • model (LightningModule | type[LightningModule] | None)

  • in_channels (int)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • dataset_kwargs (dict[str, Any] | None)

  • variant (str)

  • pretrained (bool)

  • kwargs (Any)

REQUIRED_IMAGE_SIZE: tuple[int, int] = (224, 224)
class datamint.lightning.trainers.UNETRPPTrainer(dataset=None, project=None, *, patch_crop_size=(128, 128, 128), feature_size=16, num_heads=4, depths=None, sw_overlap=0.5, in_channels=1, model=None, loss_fn=None, batch_size=1, num_workers=4, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, dataset_kwargs=None, **kwargs)

Bases: VolumeSegmentationTrainer

Convenience trainer pre-configured for UNETR++.

UNETR++ is a true 3-D segmentation model built on a hierarchical transformer encoder with Efficient Paired Attention (EPA) and a CNN decoder with skip connections.

Reference: Shaker et al., “UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation”, IEEE TMI 2024.

Example:

trainer = UNETRPPTrainer(project='CT_Liver')
results = trainer.fit()
Parameters:

  • dataset (DatamintBaseDataset | None) – A pre-built DatamintBaseDataset. Mutually exclusive with project.

  • project (str | Project | None) – Project name or Project object.

  • patch_crop_size (tuple[int, int, int]) – (D, H, W) patch randomly cropped from each volume during training and used as the sliding-window patch size at eval. Must be divisible by 32 in each dimension. Default (128, 128, 128).

  • feature_size (int) – Base channel width F. Encoder stage dims are [2F, 4F, 8F, 16F]. Default 16 matches the original paper.

  • num_heads (int) – Number of attention heads in EPA. Default 4.

  • depths (list[int] | None) – Number of transformer blocks per encoder stage (list of 4 ints). Default [3, 3, 3, 3] from the paper’s Synapse config.

  • sw_overlap (float) – Overlap ratio for sliding-window inference. Higher values reduce tiling artefacts at the cost of more compute. Default 0.5.

  • in_channels (int) – Number of input image channels (e.g. 1 for CT, 4 for multi-modal MRI). Default 1.

  • batch_size (int) – Training batch size (number of volumes). Defaults to 1 because 3-D volumes typically differ in spatial size.

  • loss_fn (Module | None) – Custom loss. Defaults to combined BCE + Dice.

  • max_epochs (int) – Maximum training epochs.

  • early_stopping_patience (int | None) – Epochs without val improvement before stop.

  • mlflow_experiment_name (str | None) – MLflow experiment name (auto-generated if None).

  • model_name (str | None) – Name for the model in the registry (auto-generated if None).

  • auto_deploy_adapter (bool) – Auto-generate a deployment adapter after training.

  • trainer_kwargs (dict[str, Any] | None) – Extra kwargs forwarded to lightning.Trainer.

  • dataset_kwargs (dict[str, Any] | None) – Extra kwargs forwarded to VolumeDataset.

  • model (LightningModule | type[LightningModule] | None)

  • num_workers (int)

  • split_as_of_timestamp (str | None)

  • kwargs (Any)

class datamint.lightning.trainers.UNetPPTrainer(dataset=None, project=None, *, image_size=None, slice_axis=None, model=None, in_channels=3, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, dataset_kwargs=None, encoder_name='resnet34', **kwargs)

Bases: SemanticSegmentation2DTrainer

Convenience trainer pre-configured for UNet++ with stronger augmentations.

Adds elastic transform and grid distortion to the default training pipeline — augmentations that are particularly effective for medical image segmentation.

Example:

trainer = UNetPPTrainer(
    project='BUS_Segmentation',
    encoder_name='resnet34',)
results = trainer.fit()
Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • image_size (int | tuple[int, int] | None)

  • slice_axis (Literal['axial', 'sagittal', 'coronal'] | int | None)

  • model (LightningModule | type[LightningModule] | None)

  • in_channels (int)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • dataset_kwargs (dict[str, Any] | None)

  • encoder_name (str)

  • kwargs (Any)

class datamint.lightning.trainers.VolumeSegmentationTrainer(*, patch_crop_size=(128, 128, 128), batch_size=1, **kwargs)

Bases: SegmentationTrainer

Abstract trainer for true 3-D volumetric segmentation.

Uses VolumeDataset directly (no slicing). Subclasses must implement _build_model().

Training uses the full volume (or a random patch cropped inside the Lightning module’s training_step). Eval uses one full volume at a time (eval_batch_size=1); sliding-window inference is handled by the Lightning module.

Parameters:

  • patch_crop_size (tuple[int, int, int]) – (D, H, W) patch size used during training. Forwarded to the Lightning module for random cropping.

  • batch_size (int) – Number of volumes per training batch. Defaults to 1 because 3-D volumes cannot be stacked unless they are all the same spatial size — use 1 unless you are certain all your volumes match.

  • to (All remaining keyword arguments are forwarded)

:param BaseTrainer.:

Example:

class MyTrainer(VolumeSegmentationTrainer):
    def _build_model(self, loss_fn, metrics):
        return MyLightningModule(...)

trainer = MyTrainer(project='CT_Liver')
results = trainer.fit()
Parameters:

kwargs (Any)

class datamint.lightning.trainers.YOLOXTrainer(dataset=None, project=None, *, model_size='s', conf_thre=0.25, nms_thre=0.45, image_size=640, model=None, loss_fn=None, batch_size=8, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=50, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, dataset_kwargs=None, **kwargs)

Bases: DetectionTrainer

One-liner trainer for anchor-free object detection using YOLOX.

Wraps YOLOX (Apache 2.0) with sensible defaults so detection training requires only a project name and, optionally, a model size:

trainer = YOLOXTrainer(project='thyroid_nodules', model_size='s')
results = trainer.fit()
Parameters:

  • dataset (DatamintBaseDataset | None) – A pre-built ImageDataset. Mutually exclusive with project.

  • project (str | Project | None) – Project name or Project object. A ImageDataset is created automatically.

  • model_size (str) – YOLOX size variant — 'nano', 'tiny', 's', 'm', 'l', or 'x'. Defaults to 's', which balances speed and accuracy for most medical imaging tasks.

  • conf_thre (float) – Combined objectness × class-confidence threshold applied during NMS at inference time.

  • nms_thre (float) – IoU threshold for non-maximum suppression.

  • image_size (int | tuple[int, int]) – Target (H, W) size or a single int for square images. Both training and evaluation images are resized to this before being fed to YOLOX.

  • batch_size (int) – Training batch size.

  • num_workers (int) – DataLoader workers.

  • train_transform (BaseCompose | None) – Custom albumentations transform for training. When None the trainer uses its built-in augmentation pipeline.

  • eval_transform (BaseCompose | None) – Custom albumentations transform for val/test. When None the trainer uses resize + normalize.

  • max_epochs (int) – Maximum training epochs.

  • early_stopping_patience (int | None) – Patience for early stopping on val/map. Set to None to disable.

  • mlflow_experiment_name (str | None) – MLflow experiment name.

  • model_name (str | None) – Name for the model in the registry.

  • split_as_of_timestamp (str | None) – Historical timestamp for reproducible splits.

  • auto_deploy_adapter (bool) – Auto-log a deploy adapter after training.

  • trainer_kwargs (dict[str, Any] | None) – Extra kwargs forwarded to lightning.Trainer.

  • dataset_kwargs (dict[str, Any] | None) – Extra kwargs forwarded to ImageDataset.

  • model (LightningModule | type[LightningModule] | None)

  • loss_fn (Module | None)

  • kwargs (Any)

BaseTrainer

Base trainer abstraction for Datamint training workflows.

Defines the BaseTrainer template that orchestrates the full pipeline: dataset → datamodule → model → Lightning Trainer → MLflow → deploy.

class datamint.lightning.trainers.base_trainer.BaseTrainer(dataset=None, project=None, *, dataset_kwargs=None, model=None, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, **kwargs)

Bases: ABC

Abstract base trainer encapsulating an end-to-end training workflow.

Subclasses provide task-specific defaults for model architecture, transforms, loss, and metrics by overriding the _build_* / _default_* hooks. Users typically only need to specify a project (or dataset) and optionally override a few settings.

Parameters:

  • dataset (DatamintBaseDataset | None) – A pre-built DatamintBaseDataset. Mutually exclusive with project.

  • project (str | Project | None) – Project name or Project object used to auto-build a dataset when dataset is None.

  • model (LightningModule | type[LightningModule] | None) – A user-provided LightningModule. When None the trainer builds a default one via _build_model().

  • loss_fn (Module | None) – Custom loss function forwarded to the default model. Ignored when model is provided (the user’s module owns its own loss).

  • batch_size (int) – Training batch size.

  • num_workers (int) – DataLoader workers.

  • train_transform (BaseCompose | None) – Albumentations transform for training. When None the trainer uses _train_transform().

  • eval_transform (BaseCompose | None) – Albumentations transform for val/test. When None the trainer uses _eval_transform().

  • split_as_of_timestamp (str | None) – Historical timestamp used to resolve project-scoped dataset splits during training. When omitted, the resolved project split datasets capture the current UTC timestamp and training lineage logs it via MLflow.

  • max_epochs (int) – Maximum number of training epochs.

  • early_stopping_patience (int | None) – Epochs without improvement before stopping. Set to None to disable early stopping.

  • mlflow_experiment_name (str | None) – MLflow experiment name. Auto-generated from the project name when None.

  • model_name (str | None) – Name for the model in the registry. Auto-generated when None.

  • auto_deploy_adapter (bool) – When True, auto-generate a DatamintModel adapter after training.

  • trainer_kwargs (dict[str, Any] | None) – Extra keyword arguments forwarded to lightning.Trainer.

  • dataset_kwargs (dict[str, Any] | None)

  • kwargs (Any)

property datamodule: DatamintDataModule
property dataset: DatamintBaseDataset
property experiment_name: str
fit()

Run the full training pipeline.

Return type:

dict[str, Any]

Returns:

Dictionary with keys 'trainer', 'model', 'test_results', and 'adapter' (when auto_deploy_adapter is enabled).

property model: LightningModule
test(register_model=True)

Run evaluation on the test split in a fresh run.

Parameters:

register_model (bool) – When True, run a zero-epoch fit first so the checkpoint callback saves the current model to MLflow and registers it after test metrics are logged.

Return type:

list[Mapping[str, float]]

Segmentation Trainers

2-D semantic segmentation trainer.

class datamint.lightning.trainers.seg2d_trainer.SemanticSegmentation2DTrainer(*, image_size=None, slice_axis=None, model=None, in_channels=3, trainer_kwargs=None, **kwargs)

Bases: SegmentationTrainer

Trainer for 2-D semantic segmentation.

Default model: UNet++ (segmentation_models_pytorch) with a resnet34 encoder pretrained on ImageNet.

When pointed at a project made of 3-D volumes, the trainer automatically converts it to a SlicedVolumeDataset and trains on 2-D slices instead.

Parameters:

  • slice_axis (Literal['axial', 'sagittal', 'coronal'] | int | None) – Slice axis override for 3-D volume projects. When omitted, the trainer tries to infer the most sensible anatomical plane and falls back to 'axial'.

  • image_size (int | tuple[int, int] | None) – Target image size (H, W) or a single int for square images. Forwarded to default transforms. When None a sensible default is chosen.

  • in_channels (int) – Number of input image channels. Defaults to 3.

  • to (All remaining keyword arguments are forwarded)

:param BaseTrainer.:

Example:

trainer = SemanticSegmentation2DTrainer(project='BUS_Segmentation')
results = trainer.fit()
Parameters:

  • model (LightningModule | type[LightningModule] | None)

  • trainer_kwargs (dict[str, Any] | None)

  • kwargs (Any)

slice_axis: Literal['axial', 'sagittal', 'coronal'] | int | None

3-D semantic segmentation trainer (slice-based).

class datamint.lightning.trainers.seg3d_trainer.SemanticSegmentation3DTrainer(*, slice_axis='axial', encoder_name='resnet34', in_channels=3, image_size=None, **kwargs)

Bases: SegmentationTrainer

Trainer for 3-D semantic segmentation via per-slice 2-D training.

Builds a VolumeDataset, slices it along the chosen axis, and trains a 2-D segmentation model on individual slices.

Parameters:

  • slice_axis (str | int) – Slicing axis — 'axial', 'sagittal', 'coronal', or an integer axis index.

  • encoder_name (str) – SMP encoder backbone.

  • in_channels (int) – Number of input channels.

  • image_size (int | tuple[int, int] | None) – Optional target image size (H, W) or a single int for square images. When omitted, training keeps the original slice size.

Example:

trainer = SemanticSegmentation3DTrainer(
    project='CT_Liver',
    slice_axis='axial',
)
results = trainer.fit()
Parameters:

kwargs (Any)

Shared base for segmentation trainers (2-D and 3-D).

class datamint.lightning.trainers.segmentation_trainer.SegmentationTrainer(dataset=None, project=None, *, dataset_kwargs=None, model=None, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, **kwargs)

Bases: BaseTrainer

Abstract trainer for segmentation tasks.

Provides shared defaults:

  • Loss – combined BCE + Dice (_BCEDiceLoss).

  • Metrics – Mean IoU and Generalised Dice Score (torchmetrics).

  • Monitor – val/iou (maximise).

Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • dataset_kwargs (dict[str, Any] | None)

  • model (LightningModule | type[LightningModule] | None)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • kwargs (Any)

Classification Trainers

Image classification trainers.

class datamint.lightning.trainers.classification_trainer.ClassificationTrainer(dataset=None, project=None, *, dataset_kwargs=None, model=None, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, **kwargs)

Bases: BaseTrainer

Abstract trainer for classification tasks.

Provides shared defaults:

  • Loss – CrossEntropyLoss.

  • Metrics – Multiclass Accuracy and macro F1 (torchmetrics).

  • Monitor – val/accuracy (maximise).

Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • dataset_kwargs (dict[str, Any] | None)

  • model (LightningModule | type[LightningModule] | None)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • kwargs (Any)

class datamint.lightning.trainers.classification_trainer.ImageClassificationTrainer(*, architecture='resnet34', pretrained=True, image_size=None, **kwargs)

Bases: ClassificationTrainer

Trainer for image classification tasks.

Default model: ResNet-34 (via timm) pretrained on ImageNet.

Parameters:

  • architecture (str) – timm model name. Defaults to 'resnet34'.

  • pretrained (bool) – Use pretrained weights. Defaults to True.

  • image_size (int | tuple[int, int] | None) – Optional target image size (H, W) or a single int for square images. When omitted, the trainer keeps the original image size instead of forcing a resize.

Example:

trainer = ImageClassificationTrainer(project='ChestXray')
results = trainer.fit()
Parameters:

kwargs (Any)

Specialized Trainers

class datamint.lightning.trainers.specialized.unetpp.UNetPPTrainer(dataset=None, project=None, *, image_size=None, slice_axis=None, model=None, in_channels=3, loss_fn=None, batch_size=16, num_workers=4, train_transform=None, eval_transform=None, split_as_of_timestamp=None, max_epochs=1, early_stopping_patience=10, mlflow_experiment_name=None, model_name=None, auto_deploy_adapter=True, trainer_kwargs=None, dataset_kwargs=None, encoder_name='resnet34', **kwargs)

Bases: SemanticSegmentation2DTrainer

Convenience trainer pre-configured for UNet++ with stronger augmentations.

Adds elastic transform and grid distortion to the default training pipeline — augmentations that are particularly effective for medical image segmentation.

Example:

trainer = UNetPPTrainer(
    project='BUS_Segmentation',
    encoder_name='resnet34',)
results = trainer.fit()
Parameters:

  • dataset (DatamintBaseDataset | None)

  • project (str | Project | None)

  • image_size (int | tuple[int, int] | None)

  • slice_axis (Literal['axial', 'sagittal', 'coronal'] | int | None)

  • model (LightningModule | type[LightningModule] | None)

  • in_channels (int)

  • loss_fn (Module | None)

  • batch_size (int)

  • num_workers (int)

  • train_transform (BaseCompose | None)

  • eval_transform (BaseCompose | None)

  • split_as_of_timestamp (str | None)

  • max_epochs (int)

  • early_stopping_patience (int | None)

  • mlflow_experiment_name (str | None)

  • model_name (str | None)

  • auto_deploy_adapter (bool)

  • trainer_kwargs (dict[str, Any] | None)

  • dataset_kwargs (dict[str, Any] | None)

  • encoder_name (str)

  • kwargs (Any)

Lightning Modules

Subclass these modules to plug custom architectures into the Datamint trainer workflow while keeping Datamint-native inference and deployment behavior.

Combined LightningModule + BaseDatamintModel base for built-in trainers.

class datamint.lightning.trainers.lightning_modules.base.DatamintLightningModule(settings=None, transform=None)

Bases: LightningModule, BaseDatamintModel

A LightningModule that is also a BaseDatamintModel.

Built-in trainers use this as the base for their default models so that the trained module can be logged once with datamint_flavor — no separate adapter step is required.

Inherits from both LightningModule and BaseDatamintModel.

Parameters:

  • settings (ModelSettings | None)

  • transform (BasicTransform | BaseCompose | None)

enable_sample_logging(enabled=True)

Enable or disable per-sample metric accumulation.

Parameters:

enabled (bool)

Return type:

None

load_context(context)

Move weights to the configured device and set eval mode on MLflow load.

Parameters:

context (PythonModelContext)

Return type:

None

mlflow_model_id: str | None
on_test_epoch_end()

Called in the test loop at the very end of the epoch.

Return type:

None

on_test_start()

Called at the beginning of testing.

Return type:

None

LightningModule wrapper for segmentation tasks.

class datamint.lightning.trainers.lightning_modules.segmentation_module.SegmentationModule(loss_fn=None, metrics_factories={}, class_names=None, transform=None, lr=0.0001)

Bases: DatamintLightningModule

Parameters:

  • loss_fn (Module | None)

  • metrics_factories (dict[str, Callable[[], Any]])

  • class_names (list[str] | None)

  • transform (BasicTransform | BaseCompose | None)

  • lr (float)

configure_optimizers()

Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you’d need one. But in the case of GANs or similar you might have multiple. Optimization with multiple optimizers only works in the manual optimization mode.

Returns:

Any of these 6 options.

  • Single optimizer.

  • List or Tuple of optimizers.

  • Two lists - The first list has multiple optimizers, and the second has multiple LR schedulers (or multiple lr_scheduler_config).

  • Dictionary, with an "optimizer" key, and (optionally) a "lr_scheduler" key whose value is a single LR scheduler or lr_scheduler_config.

  • None - Fit will run without any optimizer.

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = {
    # REQUIRED: The scheduler instance
    "scheduler": lr_scheduler,
    # The unit of the scheduler's step size, could also be 'step'.
    # 'epoch' updates the scheduler on epoch end whereas 'step'
    # updates it after a optimizer update.
    "interval": "epoch",
    # How many epochs/steps should pass between calls to
    # `scheduler.step()`. 1 corresponds to updating the learning
    # rate after every epoch/step.
    "frequency": 1,
    # Metric to monitor for schedulers like `ReduceLROnPlateau`
    "monitor": "val_loss",
    # If set to `True`, will enforce that the value specified 'monitor'
    # is available when the scheduler is updated, thus stopping
    # training if not found. If set to `False`, it will only produce a warning
    "strict": True,
    # If using the `LearningRateMonitor` callback to monitor the
    # learning rate progress, this keyword can be used to specify
    # a custom logged name
    "name": None,
}

When there are schedulers in which the .step() method is conditioned on a value, such as the torch.optim.lr_scheduler.ReduceLROnPlateau scheduler, Lightning requires that the lr_scheduler_config contains the keyword "monitor" set to the metric name that the scheduler should be conditioned on.

Metrics can be made available to monitor by simply logging it using self.log('metric_to_track', metric_val) in your LightningModule.

Note

Some things to know:

  • Lightning calls .backward() and .step() automatically in case of automatic optimization.

  • If a learning rate scheduler is specified in configure_optimizers() with key "interval" (default “epoch”) in the scheduler configuration, Lightning will call the scheduler’s .step() method automatically in case of automatic optimization.

  • If you use 16-bit precision (precision=16), Lightning will automatically handle the optimizer.

  • If you use torch.optim.LBFGS, Lightning handles the closure function automatically for you.

  • If you use multiple optimizers, you will have to switch to ‘manual optimization’ mode and step them yourself.

  • If you need to control how often the optimizer steps, override the optimizer_step() hook.

abstractmethod forward(x)

Same as torch.nn.Module.forward().

Parameters:

  • *args – Whatever you decide to pass into the forward method.

  • **kwargs – Keyword arguments are also possible.

  • x (Tensor)

Return type:

Tensor

Returns:

Your model’s output

on_test_epoch_end()

Called in the test loop at the very end of the epoch.

Return type:

None

on_train_epoch_end()

Called in the training loop at the very end of the epoch.

To access all batch outputs at the end of the epoch, you can cache step outputs as an attribute of the LightningModule and access them in this hook:

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()
Return type:

None

on_validation_epoch_end()

Called in the validation loop at the very end of the epoch.

Return type:

None

predict_image(model_input, compute_uncertainty=False, **kwargs)

Run segmentation inference, returning ImageSegmentation per resource.

Parameters:

  • compute_uncertainty (bool) – If True, also compute a predictive-entropy uncertainty score per class (see datamint.utils.uncertainty, restricted to the pixels predicted as that class) and attach it as uncertainty on each returned annotation.

  • kwargs (Any)

task_type = 'image_segmentation'

Base segmentation module for semantic segmentation tasks.

Subclasses must implement _build_model() to return the model.

Parameters:

  • in_channels – Number of input channels.

  • num_classes – Number of segmentation classes excluding background.

  • loss_fn – Loss module.

  • metrics_factories – {name: callable} where each callable returns a fresh torchmetrics.Metric. One instance is created per stage (train / val / test).

  • class_names – Human-readable label for each class.

  • image_size – (height, width) used during inference.

  • lr – Learning rate for AdamW.

test_step(batch, batch_idx)

Operates on a single batch of data from the test set. In this step you’d normally generate examples or calculate anything of interest such as accuracy.

Parameters:

  • batch (dict) – The output of your data iterable, normally a DataLoader.

  • batch_idx (int) – The index of this batch.

  • dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Return type:

Tensor

Returns:

  • Tensor - The loss tensor

  • dict - A dictionary. Can include any keys, but must include the key 'loss'.

  • None - Skip to the next batch.

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...


# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

If you pass in multiple test dataloaders, test_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    x, y = batch

    # implement your own
    out = self(x)

    if dataloader_idx == 0:
        loss = self.loss0(out, y)
    else:
        loss = self.loss1(out, y)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs separately for each dataloader
    self.log_dict({f"test_loss_{dataloader_idx}": loss, f"test_acc_{dataloader_idx}": acc})

Note

If you don’t need to test you don’t need to implement this method.

Note

When the test_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of the test epoch, the model goes back to training mode and gradients are enabled.

training_step(batch, batch_idx)

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Parameters:

  • batch (dict) – The output of your data iterable, normally a DataLoader.

  • batch_idx (int) – The index of this batch.

  • dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Return type:

Tensor

Returns:

  • Tensor - The loss tensor

  • dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.

  • None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note

When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

validation_step(batch, batch_idx)

Operates on a single batch of data from the validation set. In this step you’d might generate examples or calculate anything of interest like accuracy.

Parameters:

  • batch (dict) – The output of your data iterable, normally a DataLoader.

  • batch_idx (int) – The index of this batch.

  • dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Return type:

Tensor

Returns:

  • Tensor - The loss tensor

  • dict - A dictionary. Can include any keys, but must include the key 'loss'.

  • None - Skip to the next batch.

# if you have one val dataloader:
def validation_step(self, batch, batch_idx): ...


# if you have multiple val dataloaders:
def validation_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single validation dataset
def validation_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'val_loss': loss, 'val_acc': val_acc})

If you pass in multiple val dataloaders, validation_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple validation dataloaders
def validation_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    x, y = batch

    # implement your own
    out = self(x)

    if dataloader_idx == 0:
        loss = self.loss0(out, y)
    else:
        loss = self.loss1(out, y)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs separately for each dataloader
    self.log_dict({f"val_loss_{dataloader_idx}": loss, f"val_acc_{dataloader_idx}": acc})

Note

If you don’t need to validate you don’t need to implement this method.

Note

When the validation_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, the model goes back to training mode and gradients are enabled.

LightningModule wrapper for image classification tasks.

class datamint.lightning.trainers.lightning_modules.classification_module.ClassificationModule(model_name, num_classes, loss_fn, metrics_factories, class_names, image_size, lr=0.0001, pretrained=True, transform=None)

Bases: DatamintLightningModule

Parameters:

  • model_name (str)

  • num_classes (int)

  • loss_fn (Module)

  • metrics_factories (dict[str, Callable[[], Any]])

  • class_names (list[tuple[str, str]])

  • image_size (tuple[int, int] | None)

  • lr (float)

  • pretrained (bool)

  • transform (BasicTransform | BaseCompose | None)

configure_optimizers()

Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you’d need one. But in the case of GANs or similar you might have multiple. Optimization with multiple optimizers only works in the manual optimization mode.

Returns:

Any of these 6 options.

  • Single optimizer.

  • List or Tuple of optimizers.

  • Two lists - The first list has multiple optimizers, and the second has multiple LR schedulers (or multiple lr_scheduler_config).

  • Dictionary, with an "optimizer" key, and (optionally) a "lr_scheduler" key whose value is a single LR scheduler or lr_scheduler_config.

  • None - Fit will run without any optimizer.

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = {
    # REQUIRED: The scheduler instance
    "scheduler": lr_scheduler,
    # The unit of the scheduler's step size, could also be 'step'.
    # 'epoch' updates the scheduler on epoch end whereas 'step'
    # updates it after a optimizer update.
    "interval": "epoch",
    # How many epochs/steps should pass between calls to
    # `scheduler.step()`. 1 corresponds to updating the learning
    # rate after every epoch/step.
    "frequency": 1,
    # Metric to monitor for schedulers like `ReduceLROnPlateau`
    "monitor": "val_loss",
    # If set to `True`, will enforce that the value specified 'monitor'
    # is available when the scheduler is updated, thus stopping
    # training if not found. If set to `False`, it will only produce a warning
    "strict": True,
    # If using the `LearningRateMonitor` callback to monitor the
    # learning rate progress, this keyword can be used to specify
    # a custom logged name
    "name": None,
}

When there are schedulers in which the .step() method is conditioned on a value, such as the torch.optim.lr_scheduler.ReduceLROnPlateau scheduler, Lightning requires that the lr_scheduler_config contains the keyword "monitor" set to the metric name that the scheduler should be conditioned on.

Metrics can be made available to monitor by simply logging it using self.log('metric_to_track', metric_val) in your LightningModule.

Note

Some things to know:

  • Lightning calls .backward() and .step() automatically in case of automatic optimization.

  • If a learning rate scheduler is specified in configure_optimizers() with key "interval" (default “epoch”) in the scheduler configuration, Lightning will call the scheduler’s .step() method automatically in case of automatic optimization.

  • If you use 16-bit precision (precision=16), Lightning will automatically handle the optimizer.

  • If you use torch.optim.LBFGS, Lightning handles the closure function automatically for you.

  • If you use multiple optimizers, you will have to switch to ‘manual optimization’ mode and step them yourself.

  • If you need to control how often the optimizer steps, override the optimizer_step() hook.

forward(x)

Same as torch.nn.Module.forward().

Parameters:

  • *args – Whatever you decide to pass into the forward method.

  • **kwargs – Keyword arguments are also possible.

  • x (Tensor)

Return type:

Tensor

Returns:

Your model’s output

on_test_epoch_end()

Called in the test loop at the very end of the epoch.

Return type:

None

on_train_epoch_end()

Called in the training loop at the very end of the epoch.

To access all batch outputs at the end of the epoch, you can cache step outputs as an attribute of the LightningModule and access them in this hook:

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # free up the memory
        self.training_step_outputs.clear()
Return type:

None

on_validation_epoch_end()

Called in the validation loop at the very end of the epoch.

Return type:

None

predict_default(model_input, compute_uncertainty=False, **kwargs)

Run classification inference, returning ImageClassification per resource.

Parameters:

  • compute_uncertainty (bool) – If True, also compute a predictive-entropy uncertainty score (see datamint.utils.uncertainty) and attach it as uncertainty on each returned annotation.

  • kwargs (Any)

task_type = 'image_classification'

Generic image classification module backed by timm.

Parameters:

  • model_name – timm model name (e.g. 'resnet34', 'efficientnet_b0').

  • num_classes – Number of output classes.

  • loss_fn – Loss module.

  • metrics_factories – {name: callable} – see SegmentationModule.

  • lr – Learning rate for AdamW.

  • pretrained – Use pretrained weights.

  • image_size – Optional inference resize target (H, W). When omitted, predictions keep the original image size.

test_step(batch, batch_idx)

Operates on a single batch of data from the test set. In this step you’d normally generate examples or calculate anything of interest such as accuracy.

Parameters:

  • batch (dict) – The output of your data iterable, normally a DataLoader.

  • batch_idx (int) – The index of this batch.

  • dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Return type:

Tensor

Returns:

  • Tensor - The loss tensor

  • dict - A dictionary. Can include any keys, but must include the key 'loss'.

  • None - Skip to the next batch.

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...


# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

If you pass in multiple test dataloaders, test_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    x, y = batch

    # implement your own
    out = self(x)

    if dataloader_idx == 0:
        loss = self.loss0(out, y)
    else:
        loss = self.loss1(out, y)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs separately for each dataloader
    self.log_dict({f"test_loss_{dataloader_idx}": loss, f"test_acc_{dataloader_idx}": acc})

Note

If you don’t need to test you don’t need to implement this method.

Note

When the test_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of the test epoch, the model goes back to training mode and gradients are enabled.

training_step(batch, batch_idx)

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Parameters:

  • batch (dict) – The output of your data iterable, normally a DataLoader.

  • batch_idx (int) – The index of this batch.

  • dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Return type:

Tensor

Returns:

  • Tensor - The loss tensor

  • dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.

  • None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note

When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

validation_step(batch, batch_idx)

Operates on a single batch of data from the validation set. In this step you’d might generate examples or calculate anything of interest like accuracy.

Parameters:

  • batch (dict) – The output of your data iterable, normally a DataLoader.

  • batch_idx (int) – The index of this batch.

  • dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Return type:

Tensor

Returns:

  • Tensor - The loss tensor

  • dict - A dictionary. Can include any keys, but must include the key 'loss'.

  • None - Skip to the next batch.

# if you have one val dataloader:
def validation_step(self, batch, batch_idx): ...


# if you have multiple val dataloaders:
def validation_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single validation dataset
def validation_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'val_loss': loss, 'val_acc': val_acc})

If you pass in multiple val dataloaders, validation_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple validation dataloaders
def validation_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    x, y = batch

    # implement your own
    out = self(x)

    if dataloader_idx == 0:
        loss = self.loss0(out, y)
    else:
        loss = self.loss1(out, y)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs separately for each dataloader
    self.log_dict({f"val_loss_{dataloader_idx}": loss, f"val_acc_{dataloader_idx}": acc})

Note

If you don’t need to validate you don’t need to implement this method.

Note

When the validation_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, the model goes back to training mode and gradients are enabled.