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.

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},
)

Using an External Model Inside 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.

Related Examples

• BUSI trainer notebook

• External model deployment tutorial

• UNETR++ on the Synapse dataset

• nnU-Net on the Synapse dataset

For a side-by-side comparison of raw PyTorch / Lightning code versus Datamint, see datamint_vs_raw_pytorch.