Training¶

Set up training¶

All settings are expected to be specified in main_config.yaml.
Training settings live under the trainconfig section. Defaults, validation, and auto-derived values are managed through SkiNet.ML.configs.train_configs.train_config.TrainConfig.
precision is auto-detected from accelerator: GPU/CUDA/MPS → "16-mixed", CPU → "32-true". Override explicitly if needed.
For an Optuna sweep on GPU keep precision: "16-mixed"; for CPU sweeps either omit precision (auto-detected) or set a CPU-supported value.

Key TrainConfig fields¶

Field	Default	Description
`batch_size`	`8`	Samples per training batch
`num_workers`	auto (CPU count / DDP devices)	DataLoader workers; auto-divided among DDP processes
`pin_memory`	`True` on GPU, `False` on CPU/MPS	Auto-set from accelerator
`prefetch_factor`	`None`	Batches pre-loaded per worker; ignored when `num_workers=0`
`cache_in_ram`	`True`	Cache dataset in RAM before training; set `False` for large datasets (e.g. ISIC full split)
`use_torch_compile`	`False`	Wrap model with `torch.compile` for faster inference; first forward pass incurs JIT compilation overhead
`loss_name`	`BCE_DICE`	`BCE`, `DICE`, or `BCE_DICE` (equal 0.5/0.5 weight)
`optimizer_name`	`"adamw"`	`"adam"` or `"adamw"`
`lr`	`1e-4`	Base learning rate (linearly scaled by batch size in Optuna sweeps)
`weight_decay`	`1e-4`	L2 regularisation strength
`optimal_threshold`	`None`	`None` = sweep 51 thresholds each validation epoch; or a fixed float in [0, 1]
`max_epochs`	`1`	Maximum training epochs
`accelerator`	`"auto"`	Resolves to `gpu` / `mps` / `cpu`
`precision`	auto	Derived from accelerator; override if needed
`deterministic`	`True`	See Reproducibility
`seed`	`42`	Global RNG seed passed to `L.seed_everything`
`use_lr_scheduler`	`True`	Enable/disable the LR scheduler
`scheduler_type`	`"reduce_on_plateau"`	`"reduce_on_plateau"` or `"cosine_annealing"`
`cosine_annealing_config.T_max`	`None` (→ `max_epochs`)	Period of cosine annealing; auto-set to `max_epochs` when `None`
`cosine_annealing_config.eta_min`	`1e-6`	Minimum learning rate at the end of each cosine cycle

Lightning model¶

The core training logic lives in SkiNet.ML.model.lightning_model.LightningModel (a L.LightningModule subclass).

Metrics¶

Metric logged	Description
`train_loss` / `val_loss` / `test_loss`	Epoch-mean loss for the configured `loss_name`
`train_dice` / `val_dice` / `test_dice`	`BinaryF1Score` at the fixed or swept threshold
`train_iou` / `val_iou` / `test_iou`	`BinaryJaccardIndex`
`val_best_dice_at_threshold`	Best dataset-aggregated Dice found across the threshold sweep
`val_mean_dice_per_image`	Mean of per-image Dice at the best threshold; schema default monitor (`MetricsKey.default_monitor()`)
`val_optimal_threshold`	Threshold that achieved `val_best_dice_at_threshold`
`val_dice_threshold_gain`	Dice gain from using the optimal threshold vs. 0.5
`val_threshold_used`	Threshold actually applied during validation

Optimizer and scheduler¶

Optimizer: Adam or AdamW (set via optimizer_name), with configurable lr and weight_decay.
Scheduler: controlled by scheduler_type ("reduce_on_plateau" or "cosine_annealing"); toggled via use_lr_scheduler.
- "reduce_on_plateau" (schema default): ReduceLROnPlateau, mode "max", schema defaults patience 5 / factor 0.5 (the shipped main_config.yaml sets patience 3). It monitors the propagated SWEEP_CONFIG.monitor (schema default val_mean_dice_per_image; val_best_dice_at_threshold in the shipped config) — do not set monitor under lr_scheduler_config. Configured via lr_scheduler_config.
- "cosine_annealing": CosineAnnealingLR with T_max (defaults to max_epochs when None) and eta_min (default 1e-6). Configured via cosine_annealing_config.

Mixed precision and gradient scale monitoring¶

on_before_optimizer_step() reads trainer.precision_plugin.scaler and logs grad_scale each step. A sudden drop or collapse in grad_scale is the first diagnostic signal for mixed-precision instability.

Non-finite detection¶

Every training, validation, and test step validates inputs, logits, masks, and loss for NaN/Inf. On detection a detailed _tensor_debug_summary() is raised with the batch index and per-tensor statistics.

Model architecture¶

SkiNet uses a UNet2D architecture configured via SkiNet.ML.configs.model_configs.unet2d_config.UNet2DModelConfig (under MODEL_CONFIG in the YAML).

Key UNet2DModelConfig fields¶

Field	Default	Description
`in_channels`	`3`	Input image channels
`out_channels_layer1`	`16`	Output channels of the first encoder layer
`number_of_layers`	`5`	Total encoder layers (decoder has `number_of_layers - 1` + one final conv)
`num_output_classes`	`1`	Segmentation output classes
`kernel`	`3`	Convolution kernel size
`stride`	`2`	Downsampling stride in encoder; upsampling factor in decoder
`encoder_residual_mode`	`"he2"`	Residual block type for encoder layers
`merge_residual_mode`	`"he2"`	Residual block type for decoder merge layers
`se_reduction`	`16`	Channel reduction ratio for SE blocks (only used when `encoder_residual_mode="se"`)
`validate_forward`	`True`	Structural validation (skip key count) during forward pass
`debug_forward`	`False`	Log warnings for near-zero skip connections (GPU-expensive; keep off in production)

Encoder residual modes¶

Configured via encoder_residual_mode in MODEL_CONFIG:

Mode	Description
`"classical"`	Standard UNet encoder: Conv-BN-Act without any residual connection
`"local_refinement"`	Post-activation: downsample → refine with residual from downsampled intermediate (Oktay et al. 2020)
`"he2"` (default)	Pre-activation with 1×1 projection shortcut: BN-Act-Conv → BN-Act-Conv + P(x) (He et al. ECCV 2016)
`"se"`	Pre-activation He2 with Squeeze-and-Excitation channel attention applied before the shortcut addition (Hu et al. CVPR 2018)

Merge (decoder) residual modes¶

Configured via merge_residual_mode in MODEL_CONFIG:

Mode	Description
`"classical"`	Standard UNet decoder: upsample → concatenate skip → Conv-BN-Act without any residual connection
`"local_refinement"`	Post-activation: project-and-sum → BN-Act → Conv-BN-Act + residual (Oktay et al. 2020)
`"he1"`	Pre-activation with one refinement conv + identity shortcut (He et al. ECCV 2016)
`"he2"` (default)	Pre-activation with two refinement convs + identity shortcut (He et al. ECCV 2016)
`"attention_gate"`	Additive attention gate (Oktay et al. MIDL 2018) gates the skip connection before merge; post-merge he2 refinement

Example YAML (the production configuration):

MODEL_CONFIG:
  encoder_residual_mode: "classical"
  merge_residual_mode: "attention_gate"

Best threshold selection¶

At the end of each validation epoch, when optimal_threshold is null, SkiNet.ML.training.training_utils.find_best_threshold() sweeps 51 evenly-spaced candidate values from 1.0 down to 0.0 using torch.linspace. For every threshold it computes true positives, false positives, and false negatives in a single vectorised broadcast across the entire validation set, then derives Dice (F1) as 2·tp / (2·tp + fp + fn). The threshold with the highest Dice is selected; when multiple thresholds tie, the highest one wins because the sweep is descending and argmax returns the first occurrence.

When optimal_threshold is set to a fixed float (as in the shipped main_config.yaml, 0.5), the sweep is skipped: val_best_dice_at_threshold and val_mean_dice_per_image are computed directly at that fixed threshold and val_optimal_threshold echoes it. In both cases val_dice_threshold_gain reports the Dice difference relative to a plain 0.5 cutoff. val_best_dice_at_threshold is the metric monitored by early stopping and Optuna.

Callbacks¶

All callbacks are opt-in via boolean flags in TrainConfig. They are wired together in SkiNet.Utils.logging.logging_callbacks_setup.setup_logging_and_callbacks().

Callback	Flag	Description
`SkiNet.Utils.logging.system_metrics.SystemMetricsThreadCallback`	always on	Background thread logs CPU%, RAM%, GPU memory (allocated/reserved), and GPU utilisation every `system_metrics_interval_sec` (default 5 s)
`SkiNet.Utils.logging.throughput.ThroughputCallback`	always on	Logs `perf/samples_per_sec` and `perf/time_per_step_ms` after every training batch; primary signal for GPU saturation and DataLoader bottleneck diagnosis
`EarlyStopping`	`use_early_stopping`	Monitors the propagated `SWEEP_CONFIG.monitor` (mode `"max"`); config via `SkiNet.ML.configs.train_configs.train_config.EarlyStoppingConfig` (schema default patience 5; `main_config.yaml` uses 30)
`ModelCheckpoint`	`use_checkpoint`	Saves the best checkpoint by the propagated `SWEEP_CONFIG.monitor`; config via `SkiNet.ML.configs.train_configs.train_config.CheckpointConfig`
`MLFlowLogger`	`use_mlflow_logger`	Logs params, metrics, model summary, and artifacts; supports nested Optuna child runs; config via `SkiNet.ML.configs.train_configs.train_config.MLflowConfig`
`LearningRateMonitor`	any logger enabled	Logs LR each epoch
`SkiNet.Utils.mlops.mlflow_callbacks.MLflowTrainingArtifactsCallback`	`use_mlflow_logger`	Logs model summary at fit start; logs early-stopping state and best checkpoint as artifacts at fit end
`LitLogger`	`use_litlogger_logger`	Lightning Studio native logger; config via `SkiNet.ML.configs.train_configs.train_config.LitLoggerConfig`

Ways to start training inside a configured environment (Docker container)¶

The options below assume you are inside a configured environment (as per SkiNet’s Docker container)

Optuna hyperparameter optimisation (HPO) sweep¶

When running optuna_sweep.py, a single MLflow parent run wraps the whole study. Each trial is a nested MLflow child run. The sampler is GridSampler (exhaustive grid, not random).

The search space is declared in SkiNet.ML.configs.train_configs.sweep_config.SweepConfig and keyed by SkiNet.Utils.experiment_keys.HyperparamKey. HyperparamKey is the single source of truth: adding a new member there automatically makes SkiNet.Utils.mlops.optuna_utils.validate_search_space() require it and allows build_objective to read it.

Current members:

Member	String key	Description
`HyperparamKey.LR`	`"lr"`	Learning rate (scaled linearly by batch size: `lr * batch_size / min_batch_size`). Fixed to a single value in `main_config.yaml` after E1 LR search; set multiple values in `SWEEP_CONFIG.lr` only when searching LR for a new architecture.
`HyperparamKey.WEIGHT_DECAY`	`"weight_decay"`	L2 regularisation strength
`HyperparamKey.BATCH_SIZE`	`"batch_size"`	Samples per training batch (also rescales LR — see below)
`HyperparamKey.NUM_WORKERS`	`"num_workers"`	DataLoader worker count
`HyperparamKey.PREFETCH_FACTOR`	`"prefetch_factor"`	Batches pre-loaded per worker
`HyperparamKey.SCHEDULER_TYPE`	`"scheduler_type"`	LR scheduler per trial: `"none"` (sets `use_lr_scheduler=False`), `"cosine_annealing"`, or `"reduce_on_plateau"`

Each field is a list of GridSampler candidates for one dimension; the effective search space is the Cartesian product of the lists, and optuna_sweep.py runs the full product (n_combos = ∏ len(list)) unless --trials N caps it. Single-element lists are held constant.

Every SweepConfig field defaults to a single value, kept consistent with the SWEEP_CONFIG block in main_config.yaml (lr=[3e-4], weight_decay=[0.0], batch_size=[8], num_workers=[2], prefetch_factor=[4], scheduler_type=["none"]), so the default grid is a 1-combo no-op sweep. To search a dimension, widen its list in the YAML — in practice vary one at a time. Note num_workers and prefetch_factor are throughput knobs that do not affect model quality, and batch_size is usually swept to find the largest size that fits GPU memory rather than as a generalisation target. When batch_size is varied, lr is rescaled per trial by scale_lr (anchored to the smallest batch in the sweep), so a sampled lr always denotes the rate at the reference batch size.

SWEEP_CONFIG.monitor and SWEEP_CONFIG.direction in the YAML are the single source of truth for the optimisation objective. SkiNet.ML.configs.experiment_config.ExperimentConfig automatically propagates monitor into early_stopping_config, checkpoint_config, and lr_scheduler_config at config-load time — do not set monitor in those sub-sections. If they disagree, a ValueError is raised at startup so the mismatch is caught before any training runs. The schema default (when SWEEP_CONFIG.monitor is unset) is val_mean_dice_per_image; the shipped main_config.yaml overrides it to val_best_dice_at_threshold, which is what the callbacks below then monitor.
```
SWEEP_CONFIG:
  monitor: "val_best_dice_at_threshold"   # <-- only place to set the metric
  direction: "maximize"
  # early_stopping_config, checkpoint_config, lr_scheduler_config: omit 'monitor' — propagated automatically
```
--monitor / --direction CLI flags are optional overrides — use them only for a one-off run without editing the YAML. When omitted, optuna_sweep.py reads directly from SWEEP_CONFIG.
Similarly, the shell env vars SWEEP_MONITOR / SWEEP_DIRECTION in on_start_gpu.sh are optional overrides; leaving them unset means the YAML values are used.
MLflow run naming:
- Parent: optuna_study_{TRAIN_CONFIG.experiment_name}_{monitor}
- Child: trial_{n}_lr{lr}_wd{weight_decay}_bs{batch_size}_nw{num_workers}_pf{prefetch_factor}_sched{scheduler_type}

Example — monitor and direction taken from SWEEP_CONFIG in the YAML (recommended):

python optuna_sweep.py --config main_config.yaml

Example — one-off override without editing the YAML:

python optuna_sweep.py --config main_config.yaml --monitor val_best_dice_at_threshold --trials 10

Regular training¶

python main_run.py --config main_config.yaml

MLflow run name: {experiment_name}_seed{seed}_{timestamp}.

Multi-seed training¶

Launch a training experiment as specified in TRAIN_CONFIG, using multiple seeds. All seeds will run under one MLflow experiment

python run_seeds.py --config main_config.yaml --seeds 42 200 300

Each seed produces an independent MLflow run.

Launching training from Lightning Studio¶

The startup scripts (on_start_gpu.sh, on_start_cpu.sh) handle Docker bootstrap on Lightning Studio and dispatch to the commands above. See development.md for the full reference: available MODE values, DATASET/ENCODER_MODES/MERGE_MODES/RELEASE_GPU env vars, dry-run, and example invocations.

Training monitoring¶

MLFlow¶

When using on_start_gpu.sh / on_start_cpu.sh: MLflow is started automatically by the script — no manual step needed.

When running Python entry points directly inside the container: start the MLflow server first.

Via the setup script (recommended):

chmod +x start_mlflow.sh
./start_mlflow.sh

Or manually (SQLite backend + local artifact store):

mlflow server \
  --backend-store-uri sqlite:////workplace/SkiNet/mlflow.db \
  --default-artifact-root file:///workplace/SkiNet/mlruns \
  --host 0.0.0.0 \
  --port 5000

Open the MLflow UI in a browser on port 5000. If using Lightning Studio, tunnel via SSH:

ssh -N -L 5000:localhost:5000 ssh_connection_string_from_your_studio@ssh.lightning.ai

GPU utilisation¶

nvidia-smi dmon -s u

Reproducibility¶

How the same seed value is guaranteed¶

Three independent seed fields are read from the YAML and, in the shipped main_config.yaml, set to the same integer 100 (Pydantic schema defaults differ: split_random_seed=42, seed_value=None, seed=42):

DATA_CONFIG.split_random_seed   # used once when datasets are split
TRANSFORM_CONFIG.seed_value     # passed into the augmentation pipeline constructor
TRAIN_CONFIG.seed               # applied globally

Then in configure_reproducibility():

L.seed_everything(train_cfg.seed, workers=True)  # seeds Python, NumPy, PyTorch, and DataLoader workers
# fallback: if TRANSFORM_CONFIG.seed_value is None (and compose_kwargs has no "seed"), set it to train_cfg.seed

All three fields are set to a constant in the YAML — nothing overrides them at runtime.

Note: Toggling train_and_evaluate(..., visualize=...) (default True, which calls visualize_augmented_data on the train split) does not change model weight initialisation: visualize_augmented_data snapshots torch/NumPy/Python/CUDA RNG state and restores it in a finally block, so the RNG seen by downstream model init and training is identical with or without visualisation. No manual re-seeding is needed.

Platform notes — Deterministic mode (Apple Silicon / MPS)¶

Lightning raises MisconfigurationException if deterministic: true is used on Apple Silicon (MPS) because the MPS backend does not support deterministic mode. This is handled automatically:

# On MPS, fall back to "warn" so Lightning logs a warning but continues.
# True determinism requires CUDA/cuDNN.
if train_cfg.deterministic and torch.backends.mps.is_available():
    deterministic = "warn"
else:
    deterministic = train_cfg.deterministic

if deterministic is True and torch.cuda.is_available():
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

trainconfig:
  deterministic: true

On Ubuntu/CUDA (the primary SkiNet target) this fully enables cuDNN deterministic mode.