README.md

Trainers

All trainers are selected by class name with --trainer <ClassName> when running main.py (launched with torchrun, see BaseTrainer). Available names:

--trainer	File	Base class	Typical dataset / model
Trainer_3D	trainer_3D.py	BaseTrainer	ATLAS, BraTS3D / UNet3D
Trainer_2Dsliced	trainer_2Dsliced.py	BaseTrainer	BraTS2D / UNet2D
Trainer_2D	trainer_2D.py	BaseTrainer	QaTaCov / UNet2D
Trainer_2D_TextGuided	trainer_2D_text_guided.py	BaseTrainerText	QaTaCovTextEmb / UNet2D + CombinedSegSigLoss
Trainer_3DText	trainer_3D_text_guided.py	BaseTrainer	BraTS3DText / UNet3DText

Available Trainers

🏋🏼‍♂️ trainer_3D.py:

Trainer for 3D volume segmentation. Supports mixed precision (--mixed_precision), clips gradients (norm 2) and calls train_sampler.set_epoch(epoch) for DDP shuffling.

Key Methods:

• _train_epoch(epoch): Iterates through patch-based training loader, computes loss, backpropagates, updates metrics
• eval_epoch(epoch, phase): Performs patch-based inference on full volumes using GridSampler/GridAggregator (hann overlap mode), computes validation/test metrics
• _inference_sampler(sample): Creates GridSampler for dense patch-based inference (uses dataset.patch_size and dataset.grid_overlap; volumes smaller than the patch are padded with CropOrPad)
• _results_dict(phase, epoch): Retrieves metrics for the epoch

🏋🏼‍♂️ trainer_2Dsliced.py:

Trainer for 2D slice-based segmentation of 3D volumes (datasets inheriting from BaseDataset2DSliced). At evaluation, predictions are argmax'ed before computing metrics.

🏋🏼‍♂️ trainer_2D.py:

Trainer for full 2D image segmentation (e.g., QaTaCov chest X-rays) without patch-based inference.

🏋🏼‍♂️ trainer_2D_text_guided.py:

Trainer_2D_TextGuided: report-guided 2D trainer built on BaseTrainerText.

• Training: calls model(image, return_bottleneck=True), projects bottleneck and the sample's text_emb (precomputed embedding, [B, 768]) with the guidance head, and computes the loss with img_emb/txt_emb (e.g. CombinedSegSigLoss). A KeyError is raised if a training sample has no text_emb. With use_text_guidance=False it behaves like Trainer_2D
• Evaluation: segmentation only (no text needed), so val/test samples can come without embeddings
• Logs parameter counts and a gradient-flow summary at the first epoch/iteration; clips gradients of model and head (norm 1)
• Fine-tuning: set pretrained_unet_path in the config (see BaseTrainerText)

🏋🏼‍♂️ trainer_3D_text_guided.py:

Trainer_3DText: report-guided 3D trainer, to be used with BraTS3DText and UNet3DText.

• Constructor: extra pretrained_path argument (from main.py --pretrained_path); if set, the weights (raw state_dict or checkpoint with a "model" key) are loaded with strict=False after initialization
• Datasets: _init_dataset() builds a training loader of subjects without reports (get_loader('train')) and a second one of subjects with reports (get_loader('train', report=True))
• Report interleaving: after every 2 segmentation-only training iterations, one extra step is done on a batch from the (cycled) report loader, calling model(image, report) → (prediction, z_img, z_txt) and loss(prediction, label, z_img, z_txt)
• Evaluation: same patch-based GridSampler/GridAggregator inference as Trainer_3D, segmentation only
• Supports mixed precision (--mixed_precision)

Implementing a Custom Trainer

1. Create a new file in trainer/ with your trainer class that inherits from BaseTrainer (or BaseTrainerText for report guidance)
2. Implement _train_epoch(epoch): Define training loop, iterate through batches, compute loss, backpropagate, and return metrics dict
3. Implement eval_epoch(epoch, phase): Define validation/test loop, compute predictions and metrics, and return metrics dict. The returned dict must contain results[<first metric key>]['<key>_<eval_metric_type>'], which is used for best-model selection
4. Add your trainer class to __init__.py in the trainer/ folder:

   from .YourTrainer import YourTrainer

5. Use it by specifying --trainer YourTrainer when running main.py.