Skip to content

training.md

Latest commit

 

History

History
144 lines (107 loc) · 8.96 KB

training.md

File metadata and controls

144 lines (107 loc) · 8.96 KB
Raw

Segmentation Model Training

Training is launched via main.py. It loads a JSON configuration, instantiates the trainer class named by --trainer, and calls its train() method. All trainers inherit from BaseTrainer, which handles distributed setup, checkpointing, W&B logging, mixed precision and the train / validation / test loop.

Requirements

  • One or more CUDA GPUs on a single node. BaseTrainer always wraps the model in DistributedDataParallel and initializes the process group with the nccl backend, so CPU-only training is not supported.
  • Launch with torchrun. BaseTrainer.__init__ checks that RANK, LOCAL_RANK and WORLD_SIZE are set and raises a RuntimeError otherwise, so running python main.py ... directly fails.
  • Multi-node training is not supported: if the world size is larger than the number of local GPUs, the trainer raises RuntimeError.
# Single GPU (DDP is still used)
torchrun --nproc_per_node=1 main.py [arguments]

# Multi GPU, single node
torchrun --nproc_per_node=NUM_GPUS main.py [arguments]

Command Line Arguments

All arguments are defined in parse_args() in main.py.

Argument Type / default Description
--config path, required Path to the configuration JSON file (e.g. config/config_brats3d.json).
--epochs int, required Total number of training epochs.
--save_path path, required Directory for checkpoints and metric CSVs. It is created if it does not exist.
--trainer str, required Trainer class name, resolved with getattr(trainer, name) (see Trainers).
--validation flag Run validation during training. Without it model_best.pth is never written.
--val_every int, default 1 Run validation every N epochs.
--resume flag Resume from <save_path>/model_last.pth, including optimizer, scheduler, best metric, metric CSVs and W&B run id.
--debug flag Print per-epoch result dictionaries.
--eval_metric_type mean (default) or aggregated_mean Column used to select model_best.pth: <key>_mean (mean over classes) or <key>_aggregated_mean (mean over the aggregated_regions, e.g. BraTS ET/TC/WT). <key> is the key of the first entry in config.metrics. aggregated_mean requires aggregated_regions in the config.
--wandb flag Enable Weights & Biases logging (see Weights & Biases).
--save_visualizations flag Stored on the trainer as self.save_visualizations, so trainers can save visualizations during validation.
--mixed_precision fp16, bf16, default none (FP32) Enable automatic mixed precision (see Mixed precision).
--pretrained_path path, default none Weights to fine-tune from. Only Trainer_3DText uses it. The other trainers accept it and ignore it.
--seed int, default 42 Seed for random, numpy and torch (CPU and CUDA).

Pretrained weights

There are two ways to load pretrained weights, and which one applies depends on the trainer:

  • Trainer_3DText reads --pretrained_path. It loads the file after the base initialization, taking the "model" entry if the file is a trainer checkpoint and otherwise treating it as a raw state dict, with strict=False.
  • Trainer_2D_TextGuided (through BaseTrainerText) ignores --pretrained_path. It reads the top-level config key pretrained_unet_path and loads that file the same way ("model" entry or raw state dict, strict=False). If the key is missing, it trains from scratch.

Trainers

The trainers are exported from trainer/__init__.py. Pass the class name to --trainer:

--trainer File Base class Typical config
Trainer_3D trainer_3D.py BaseTrainer config_atlas.json (ATLAS), config_brats3d.json (BraTS3D, UNet3D)
Trainer_2Dsliced trainer_2Dsliced.py BaseTrainer config_brats2d.json (BraTS2D slices, UNet2D)
Trainer_2D trainer_2D.py BaseTrainer config_qatacov2d.json (QaTaCov, UNet2D)
Trainer_2D_TextGuided trainer_2D_text_guided.py BaseTrainerText config_qatacov2d_textemb.json (QaTaCovTextEmb, CombinedSegSigLoss)
Trainer_3DText trainer_3D_text_guided.py BaseTrainer No dedicated config ships with the repo. Start from config_brats3d.json and set dataset.name to BraTS3DText, add dataset.report_folder, and set model.name to UNet3DText. The trainer builds a second loader of report-paired subjects with get_loader('train', report=True).

The shipped configs contain absolute cluster paths in dataset.path, dataset.split_file and dataset.transforms. Change them before you train. All configuration keys are listed in config/README.md.

Weights & Biases

Set these environment variables before launching with --wandb:

export WANDB_ENTITY="your_entity"    # default: None (your W&B default entity)
export WANDB_PROJECT="your_project"  # default: "medical-segmentation"
  • The run name is config.name. The run config contains the JSON config plus the trainer name, eval_metric_type, mixed_precision, epochs and save_path.
  • Only rank 0 initializes W&B and logs to it. Per-epoch train, val and test metrics go under the train/, val/ and test/ prefixes.
  • The run id is saved in the checkpoint. With --resume --wandb, the same run is resumed (resume="must"). If the checkpoint has no id, a new run is started.
  • wandb.init is called with mode="offline", so runs are written locally (under ./wandb/) and are not streamed. Upload them afterwards with wandb sync.

Mixed Precision

--mixed_precision fp16 or --mixed_precision bf16 sets use_amp, amp_dtype and (for fp16 only) a GradScaler on the trainer. Without the flag, training runs in FP32.

Only the training steps of Trainer_3D and Trainer_3DText wrap the forward pass in autocast and use the scaler. The 2D trainers (Trainer_2D, Trainer_2Dsliced, Trainer_2D_TextGuided) accept the flag but train in FP32.

Checkpoints and Outputs

Everything is written to --save_path by rank 0:

File When
model_last.pth Written after every epoch, overwriting the previous one. --resume reads this file.
model_best.pth Written when the validation metric (<first metric key>_<eval_metric_type>) improves on the best value so far. The comparison is "higher is better", so put a higher-is-better metric such as DSC first in config.metrics. This file only exists when --validation is used.
train_metrics.csv, val_metrics.csv, test_metrics.csv Per-epoch metric tables (see testing.md).

A checkpoint is a dict with name, config, model (the state dict of the unwrapped model), optimizer, lr_scheduler, epoch, best_metric and wandb_run_id. Text-guided trainers based on BaseTrainerText also store the guidance-head weights and dimensions.

The test split is evaluated automatically at the last epoch. See testing.md.

Example

Train the 3D BraTS model on 4 GPUs with validation every 2 epochs, bf16, and W&B:

source .venv/bin/activate
export WANDB_ENTITY="name_of_your_entity"
export WANDB_PROJECT="name_of_your_project"

torchrun --nproc_per_node=4 main.py \
  --config config/config_brats3d.json \
  --epochs 300 \
  --save_path /path/to/runs/brats3d_unet \
  --trainer Trainer_3D \
  --validation \
  --val_every 2 \
  --eval_metric_type aggregated_mean \
  --mixed_precision bf16 \
  --wandb

To continue an interrupted run, repeat the same command with --resume added. Training restarts from the epoch after the one stored in model_last.pth.

Fine-tune the report-guided 3D model from a baseline checkpoint:

torchrun --nproc_per_node=1 main.py \
  --config config/config_brats3d_text.json \
  --epochs 100 \
  --save_path /path/to/runs/brats3d_text \
  --trainer Trainer_3DText \
  --validation \
  --pretrained_path /path/to/runs/brats3d_unet/model_best.pth

(config_brats3d_text.json is a config you create yourself, as described in Trainers.)

Related Documentation