Training Configs

pydantic model SkiNet.ML.configs.train_configs.train_config.TrainConfig[source]

Bases: BaseModel

Configuration for training.

Fields:
field accelerator: str = 'auto'

Lightning accelerator: ‘auto’, ‘gpu’, ‘mps’, or ‘cpu’.

field batch_size: int = 8

Samples per batch.

Constraints:

  • ge = 1

field cache_in_ram: bool = True

Pre-load all images into RAM at startup.

field check_val_every_n_epoch: int = 1

Validation frequency in epochs.

Constraints:

  • ge = 1

field checkpoint_config: CheckpointConfig [Optional]

Model-checkpoint callback config. Its ‘monitor’ is propagated from SWEEP_CONFIG.

field cosine_annealing_config: CosineAnnealingConfig [Optional]

CosineAnnealingLR config; used when scheduler_type=’cosine_annealing’.

field deterministic: bool = True

Enable cuDNN deterministic mode.

field devices: str | int = 'auto'

Number of devices, or ‘auto’.

field early_stopping_config: EarlyStoppingConfig [Optional]

Early-stopping callback config. Its ‘monitor’ is propagated from SWEEP_CONFIG.

field experiment_name: str = 'unet2d_experiment'

MLflow experiment name and run-name prefix. Drives run naming (run = ‘{experiment_name}_seed{seed}_{YYYYMMDD-HHMMSS}’); checkpoints are written under ‘{log_dir}/checkpoints/{run_name}’.

field litlogger_config: LitLoggerConfig [Optional]

Lightning Studio LitLogger config.

field log_dir: str = 'experiment_logs'

Local directory for logs and checkpoints.

field log_every_n_steps: int = 1

Logging frequency in steps.

Constraints:

  • ge = 1

field loss_name: LossFunctionKey = LossFunctionKey.BCE_DICE

Loss function name. Supported: ‘bce’, ‘dice’, ‘bce_dice’.

field lr: float = 0.0001

Base learning rate.

Constraints:

  • gt = 0

field lr_scheduler_config: ReduceOnPlateauConfig [Optional]

ReduceLROnPlateau config; used when scheduler_type=’reduce_on_plateau’. Its ‘monitor’ is propagated from SWEEP_CONFIG.

field max_epochs: int = 1

Maximum number of training epochs.

Constraints:

  • ge = 1

field mlflow_config: MLflowConfig [Optional]

MLflow logger config.

field num_sanity_val_steps: int = 0

Sanity validation steps before training.

Constraints:

  • ge = 0

field num_workers: int | None = None

DataLoader worker processes. When None, auto-set to os.cpu_count() (single GPU) or cpu_count // devices under DDP.

Constraints:

  • ge = 0

field optimal_threshold: float | None = None

Fixed sigmoid threshold to use instead of sweeping. When None (default), the threshold is found via grid search each validation epoch.

Constraints:

  • ge = 0.0

  • le = 1.0

field optimizer_name: str = 'adamw'

Optimizer name: ‘adam’ or ‘adamw’.

field pin_memory: bool | None = None

Pin host memory for faster H2D copies. When None, auto-set True on CUDA/GPU and False on MPS/CPU.

field precision: Literal['16-mixed', 'bf16-mixed', '32-true', '16-true', 'bf16-true'] | None = None

Training precision. When None, auto-set to ‘16-mixed’ on GPU/MPS and ‘32-true’ on CPU.

field prefetch_factor: int | None = None

Batches pre-loaded per worker. Ignored (forced None) when num_workers=0.

Constraints:

  • ge = 1

field run_test_after_fit: bool = False

Run test-set evaluation after training completes.

field scheduler_type: Literal['reduce_on_plateau', 'cosine_annealing'] = 'reduce_on_plateau'

Which scheduler sub-config is used: ‘reduce_on_plateau’ → lr_scheduler_config, ‘cosine_annealing’ → cosine_annealing_config.

field seed: int = 42

Global RNG seed passed to L.seed_everything.

Constraints:

  • ge = 0

field strategy: str = 'auto'

Lightning strategy: ‘auto’, ‘ddp’, etc.

field system_metrics_interval_sec: float = 5.0

System-metrics logging interval in seconds.

Constraints:

  • gt = 0

field test_on_val_split: bool = False

Use the validation split as the test set. Metrics may be optimistic; leave False unless you have no held-out test set.

field torch_compile_backend: str = 'inductor'

torch.compile backend. Use ‘eager’ to avoid an nvcc dependency.

field use_checkpoint: bool = False

Enable model checkpointing.

field use_early_stopping: bool = False

Enable early stopping.

field use_litlogger_logger: bool = False

Enable the Lightning Studio LitLogger.

field use_lr_scheduler: bool = True

Master toggle; when False no LR scheduler is attached.

field use_mlflow_logger: bool = False

Enable MLflow logging. Requires mlflow_config.tracking_uri to be set.

field use_torch_compile: bool = False

Wrap the model with torch.compile.

field weight_decay: float = 0.0001

L2 regularisation (weight decay).

Constraints:

  • ge = 0

validator require_tracking_uri_if_enabled  »  all fields[source]
validator set_num_workers_auto  »  all fields[source]

Auto-detect num_workers from os.cpu_count(), DDP-aware: when devices is an int > 1 (one process per device), divide the CPU budget among them to avoid oversubscription.

validator set_pin_memory_from_accelerator  »  all fields[source]

Auto-set pin_memory: True iff effective accelerator is CUDA/GPU. pin_memory is a no-op (or unsupported) for MPS/CPU.

validator set_precision_from_accelerator  »  all fields[source]

Set precision format based on the currently used device type

validator validate_prefetch_factor  »  all fields[source]
validator warn_if_testing_on_val  »  all fields[source]
pydantic model SkiNet.ML.configs.train_configs.train_config.ReduceOnPlateauConfig[source]

Bases: BaseModel

Learning rate ReduceOnPlateau scheduler configuration for PyTorch Lightning.

Fields:
field factor: float = 0.5

Multiplicative LR reduction factor.

Constraints:

  • gt = 0

  • lt = 1

field mode: Literal['min', 'max'] = 'max'

‘max’ or ‘min’ for the monitored metric.

field monitor: MetricsKey = MetricsKey.VAL_MEAN_DICE_PER_IMAGE

Metric to watch; propagated from SWEEP_CONFIG — do not set in YAML.

field patience: int = 5

Epochs without improvement before reducing the LR.

Constraints:

  • ge = 0

pydantic model SkiNet.ML.configs.train_configs.train_config.CosineAnnealingConfig[source]

Bases: BaseModel

CosineAnnealingLR scheduler configuration. T_max is set to max_epochs at runtime when None.

Fields:
field T_max: int | None = None

Cosine period in epochs; resolved to max_epochs at runtime when None.

Constraints:

  • ge = 1

field eta_min: float = 1e-06

Minimum LR at the end of each cosine cycle.

Constraints:

  • ge = 0

pydantic model SkiNet.ML.configs.train_configs.train_config.CheckpointConfig[source]

Bases: BaseModel

Configuration for model checkpointing.

Fields:
field filename: str = 'epoch{epoch:03d}'

Checkpoint filename template.

field mode: Literal['min', 'max'] = 'max'

‘max’ or ‘min’ for the monitored metric.

field monitor: MetricsKey = MetricsKey.VAL_MEAN_DICE_PER_IMAGE

Metric to watch; propagated from SWEEP_CONFIG — do not set in YAML.

field save_last: bool = True

Always save the last checkpoint.

field save_top_k: int = 1

Number of best checkpoints to keep.

Constraints:

  • ge = 0

pydantic model SkiNet.ML.configs.train_configs.train_config.EarlyStoppingConfig[source]

Bases: BaseModel

Configuration for early stopping.

Fields:
field check_finite: bool = True

Stop if the metric becomes NaN or Inf.

field divergence_threshold: float | None = None

Stop immediately if the metric falls below this value.

field min_delta: float = 0.0

Minimum change to count as an improvement.

field mode: Literal['min', 'max'] = 'max'

‘max’ or ‘min’ for the monitored metric.

field monitor: MetricsKey = MetricsKey.VAL_MEAN_DICE_PER_IMAGE

Metric to watch; propagated from SWEEP_CONFIG — do not set in YAML.

field patience: int = 5

Epochs without improvement before stopping.

Constraints:

  • ge = 0

field stopping_threshold: float | None = None

Stop immediately once the metric passes this value.

field strict: bool = True

Raise an error if the monitored metric is missing.

validator warn_monitor_is_default  »  all fields[source]
pydantic model SkiNet.ML.configs.train_configs.train_config.MLflowConfig[source]

Bases: BaseModel

Configuration for MLflow logging.

Fields:
field fallback_to_local_mlflow: bool = False

Fall back to local MLflow when the remote server is unreachable.

field log_model: bool | Literal['all'] = 'all'

Log model artifacts; ‘all’ logs every checkpoint.

field log_model_summary: bool = True

Log the model summary as an artifact.

field tracking_uri: str | None = None

MLflow tracking server URI, e.g. ‘http://127.0.0.1:5000’.

pydantic model SkiNet.ML.configs.train_configs.sweep_config.SweepConfig[source]

Bases: BaseModel

Configuration for the optuna HPO sweep.

monitor and direction are the single source of truth for the optimisation objective. ExperimentConfig propagates monitor into trainconfig.early_stopping_config, trainconfig.checkpoint_config, and trainconfig.lr_scheduler_config so all four fields are always in sync without manual YAML duplication.

Fields:
field batch_size: list[int] [Optional]

Batch-size candidates; also rescales the LR via scale_lr.

field direction: str = 'maximize'

Optuna optimisation direction: ‘maximize’ or ‘minimize’.

Constraints:

  • pattern = ^(maximize|minimize)$

field experiment_name: str = 'optuna_sweep'

MLflow experiment name for the sweep.

field lr: list[float] [Optional]

Learning-rate candidates (GridSampler dimension).

field monitor: MetricsKey = MetricsKey.VAL_MEAN_DICE_PER_IMAGE

Metric to optimise; the single source of truth, propagated into the early-stopping, checkpoint, and LR-scheduler configs.

field num_workers: list[int] [Optional]

DataLoader worker-count candidates (GridSampler dimension).

field prefetch_factor: list[int] [Optional]

Batches pre-loaded per worker (GridSampler dimension).

field scheduler_type: list[str] [Optional]

LR scheduler variants to sweep. Use ‘none’ to disable the scheduler for a given trial; ‘cosine_annealing’ and ‘reduce_on_plateau’ map directly to TrainConfig.scheduler_type.

field weight_decay: list[float] [Optional]

Weight-decay candidates (GridSampler dimension).