Training your Model

Intermediate

The trainer layer in datamint.lightning 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

The one-line trainers below come with a model already wired in — construct one with project= (plus any architecture-specific keyword arguments) and call fit():

Trainer

Task

Model

Supported format

UNetPPTrainer

2-D segmentation

UNet++ (SMP, resnet34 encoder)

2-D images; auto-slices 3-D volume projects

DeepLabV3PlusTrainer

2-D segmentation

DeepLabV3+ (SMP, ASPP)

2-D images; auto-slices 3-D volume projects

TransUNetTrainer

2-D segmentation

TransUNet (ViT-R50 hybrid encoder + CUP decoder)

2-D images, fixed 224×224; auto-slices 3-D volume projects

UNETRPPTrainer

3-D segmentation

UNETR++ (transformer encoder with EPA + CNN decoder)

3-D volumes only (true 3-D, no slicing)

NNUNetTrainer

2-D/3-D segmentation

nnU-Net v2

3-D volume projects (NIfTI/DICOM series); configuration= picks 2d, 3d_fullres, 3d_lowres, or 3d_cascade_fullres

ImageClassificationTrainer

2-D classification

Configurable timm backbone

2-D images

EfficientNetV2Trainer

2-D classification

EfficientNetV2 (timm, S/M/L/XL/rw_t)

2-D images

YOLOXTrainer

2-D object detection

YOLOX (nano/tiny/s/m/l/x)

2-D images

Note

A note on nnU-Net

NNUNetTrainer is the odd one out among the one-line trainers:

  • It requires the optional nnunet extra: pip install datamint[nnunet] (installs nnunetv2 and filelock).

  • It needs a project made of 3-D volumes, and runs nnU-Net’s own fingerprinting → planning → preprocessing → training pipeline instead of Datamint’s Lightning model/loss/metrics wiring used by every other trainer.

  • It does not accept model= — the external-model patterns described below don’t apply to it.

  • fit() can be long-running, since fingerprinting, preprocessing, and the full nnU-Net training loop all happen inside the call.

If you need more control over the model, transforms, or loss, drop down to the generic task trainers instead — they require passing model= explicitly:

SemanticSegmentation2DTrainer

A more explicit 2-D segmentation trainer when you want to control the model, transforms, or loss. It auto-detects whether the project contains 2-D images or 3-D volumes and accepts slice_axis= to override the inferred plane for volume projects.

SemanticSegmentation3DTrainer

Slice-based semantic segmentation for projects of 3-D volumes.

VolumeSegmentationTrainer

True 3-D segmentation without slicing, for volume-native architectures.

ClassificationTrainer

Abstract base for classification tasks when you want to bring your own model.

DetectionTrainer

Abstract base for object detection tasks when you want to bring your own model.

Quick Start

from datamint.lightning import UNetPPTrainer

trainer = UNetPPTrainer(
    project="BUSI_Segmentation",
    image_size=256,
    batch_size=16,
    max_epochs=20,
    accelerator="auto",
)

results = trainer.fit()
print(results["test_results"])

The built-in trainer configures the dataset, datamodule, model, MLflow logger, checkpointing, and evaluation loop for you. After fit(), the resolved objects are also available as trainer.dataset, trainer.datamodule, and trainer.model.

Inputs, Splits, and Outputs

Each trainer accepts exactly one of:

  • project=... to let Datamint build the dataset automatically, or

  • dataset=... to reuse a dataset you already configured yourself.

For SemanticSegmentation2DTrainer and UNetPPTrainer, project-backed dataset resolution is automatic: pure 2-D image projects use ImageDataset, while pure 3-D volume projects are converted to SlicedVolumeDataset. The slice plane is inferred from volume spacing and shape when possible, and falls back to 'axial'. To force a plane, pass slice_axis='coronal' or another supported axis when constructing the trainer.

When you train from a project, the trainer expects train/val/test split assignments to exist for that project. If you need strict split reproducibility across runs, pass the historical split snapshot timestamp through split_as_of_timestamp.

trainer = UNetPPTrainer(
    project="BUSI_Segmentation",
    split_as_of_timestamp="2026-04-21T12:34:56Z",
)

fit() returns a dictionary with these keys:

model

The trained model instance.

test_results

The metrics returned by Lightning test().

If you only want evaluation, use test() instead:

test_metrics = trainer.test(register_model=False)

With register_model=True, the trainer logs and registers the current model.

Resuming a Paused Run

If training is interrupted (e.g. Ctrl+C or SIGTERM) before it finishes, Lightning saves a resumable checkpoint automatically and the run’s MLflow run_id is logged:

Training paused (run_id=abcd1234...). Resume by passing resume_from='abcd1234...'.

To resume, pass that run_id (or a full checkpoint path) as resume_from when constructing a new trainer.

trainer = UNetPPTrainer(project="BUSI_Segmentation", resume_from="abcd1234...")
results = trainer.fit()

Not supported by NNUNetTrainer, which manages its own checkpointing/resuming — use NNUNetTrainer(continue_training=True) instead.

Passing Lightning Trainer Options

Any extra keyword arguments that are not consumed by Datamint are forwarded to lightning.Trainer.

trainer = UNetPPTrainer(
    project="BUSI_Segmentation",
    max_epochs=12,
    accelerator="gpu",
    devices=1,
    precision="16-mixed",
    log_every_n_steps=10,
    trainer_kwargs={"enable_progress_bar": True},
)

Training an External Model Through a Datamint Trainer

There are two supported patterns, and they are not equivalent.

Preferred: Subclass a Datamint Lightning Module

If you want to swap the network architecture but keep Datamint’s loss wiring, metrics, MLflow model behaviour, and deployment-friendly prediction methods, subclass SegmentationModule or ClassificationModule and pass the class object to model=.

import segmentation_models_pytorch as smp
from datamint.lightning import SemanticSegmentation2DTrainer
from datamint.lightning.trainers.lightning_modules import SegmentationModule


class MAnetModule(SegmentationModule):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, class_names=["benign", "malignant"], **kwargs)
        self.model = smp.MAnet(
            encoder_name="resnet50",
            encoder_weights="imagenet",
            in_channels=3,
            classes=2,
        )

    def forward(self, x):
        return self.model(x)


trainer = SemanticSegmentation2DTrainer(
    project="BUSI_Segmentation",
    image_size=256,
    model=MAnetModule,
)

results = trainer.fit()

When you pass the class object instead of an instance, the trainer instantiates it and injects task defaults through loss_fn= and metrics_factories=. This is the easiest way to plug an external architecture into the Datamint trainer workflow while keeping the resulting model Datamint-compatible.

Fully Custom LightningModule

If you already have a plain lightning.LightningModule, pass the instance through model= and implement the full training logic yourself.

import lightning as L
import segmentation_models_pytorch as smp
import torch
from datamint.lightning import SemanticSegmentation2DTrainer


class ExternalSegmentationModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = smp.Unet(
            encoder_name="resnet34",
            encoder_weights="imagenet",
            in_channels=3,
            classes=2,
        )
        self.loss_fn = torch.nn.BCEWithLogitsLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        images = batch["image"]
        masks = batch["segmentations"][:, 1:].float()
        loss = self.loss_fn(self(images), masks)
        self.log("train/loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        images = batch["image"]
        masks = batch["segmentations"][:, 1:].float()
        loss = self.loss_fn(self(images), masks)
        self.log("val/loss", loss)
        return loss

    def test_step(self, batch, batch_idx):
        images = batch["image"]
        masks = batch["segmentations"][:, 1:].float()
        loss = self.loss_fn(self(images), masks)
        self.log("test/loss", loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-4)


trainer = SemanticSegmentation2DTrainer(
    project="BUSI_Segmentation",
    model=ExternalSegmentationModule(),
    max_epochs=5,
)

This gives you the Datamint dataset, split handling, MLflow logger, and checkpointing, but the model itself remains a plain Lightning module. That means Datamint-native inference and deployment behaviour is not added automatically. If you want the trained artifact to behave like a Datamint model, prefer the SegmentationModule / ClassificationModule route, or wrap the final model in a DatamintModel afterwards.

Note

Class vs. Instance

model=MyModule and model=MyModule() are different:

  • Pass the class when you want the trainer to inject loss_fn and metrics_factories.

  • Pass an instance when the module is already fully configured and owns its entire training logic.

Caution

Segmentation batches expose masks in batch["segmentations"] and include the background channel at index 0.

Benchmarking Several Trainers

Benchmark runs several existing trainers sequentially against one shared dataset and split, then returns a ranked leaderboard.

from datamint import ImageDataset
from datamint.lightning import Benchmark
from datamint.lightning.trainers import UNetPPTrainer, DeepLabV3PlusTrainer, TransUNetTrainer

dataset = ImageDataset(project="BUSI_Segmentation", return_segmentations=True)

bench = Benchmark(
    dataset=dataset,
    trainers=[
        (UNetPPTrainer, {"encoder_name": "resnet34", "model_name": "unetpp_r34"}),
        (DeepLabV3PlusTrainer, {"model_name": "deeplabv3plus"}),
        (TransUNetTrainer, {"model_name": "transunet"}),
    ],
    main_metric="dice",
    main_metric_mode="max",
    max_epochs=20,  # shared kwarg, forwarded to every trainer
)

leaderboard = bench.run()

bench.save_config("benchmark.yaml")
# later, possibly in a different process/session:
# Benchmark.load_from_file("benchmark.yaml", dataset=dataset).run()

leaderboard is a pandas.DataFrame, one row per trainer, sorted by test/{main_metric} (falling back to val/{main_metric} if the test column is entirely missing), with columns model_name, trainer_class, val/{main_metric}, test/{main_metric}, run_id, and registered_version – so a specific trained model stays unambiguously loadable afterward via runs:/<run_id>/model.

A few rules Benchmark enforces upfront (before training anything):

  • All trainers must share one task-family ancestor: classification, 2-D segmentation, volume segmentation, or detection. A trainer outside those families (e.g. NNUNetTrainer) can only be benchmarked against other instances of that same class.

  • Every trainer spec needs an explicit, unique model_name – this keeps each trainer’s registered model separate in the MLflow Model Registry instead of silently stacking versions under one shared name.

  • The dataset’s project must already have split assignments. Benchmark pins one split_as_of_timestamp at construction (or an explicit override) and forwards it to every trainer, so all of them resolve the identical split assignments even if someone edits them mid-benchmark.

save_config()/load_from_file() round-trip the benchmark definition through YAML, independent of any live dataset object. The dataset itself is never serialized: load_from_file(path, dataset=...) requires a freshly-built dataset, mirroring the constructor’s own requirement.

Note

Have a model that was trained entirely outside Datamint and want to integrate, log, and deploy it for inference through the UI? See Bringing an External Model into Datamint instead.