Lightning API

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

Lightning Components

• DatamintDataModule
• Trainers
• Lightning Modules

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, register_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.

• register_model_name (str | None) – Name for MLflow Model 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, register_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)

• register_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, register_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)

• register_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.ImageClassificationTrainer(*, model_name='resnet34', pretrained=True, image_size=None, **kwargs)

Bases: ClassificationTrainer

Trainer for image classification tasks.

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

Parameters:

• model_name (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.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, register_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)

• register_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, register_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)

• register_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.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, register_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)

• register_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)

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, register_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.

• register_model_name (str | None) – Name for MLflow Model 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, register_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)

• register_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, register_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)

• register_model_name (str | None)

• auto_deploy_adapter (bool)

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

• kwargs (Any)

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

Bases: ClassificationTrainer

Trainer for image classification tasks.

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

Parameters:

• model_name (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, register_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)

• register_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, **kwargs)

Run segmentation inference, returning ImageSegmentation per resource.

Parameters:

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[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, **kwargs)

Run classification inference, returning ImageClassification per resource.

Parameters:

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.