Lightning Model

class SkiNet.ML.model.lightning_model.LightningModel(*args: Any, **kwargs: Any)

Bases: LightningModule

configure_optimizers() → lightning.pytorch.utilities.types.OptimizerLRScheduler

Build the optimizer and optionally a scheduler.

Scheduler type is controlled by scheduler_type: - "reduce_on_plateau": ReduceLROnPlateau (requires a monitor metric) - "cosine_annealing": CosineAnnealingLR (T_max defaults to trainer.max_epochs)

The scheduler is omitted entirely when use_lr_scheduler=False.

Returns:

optimizer alone, or Lightning-compatible dict with “optimizer” and “lr_scheduler” keys when the scheduler is enabled

Raises:

• ValueError – if optimizer_name is not “adam” or “adamw”

• ValueError – if scheduler_type is not recognised

forward(x: torch.Tensor) → torch.Tensor

Run the backbone and return raw logits (pre-sigmoid).

on_before_optimizer_step(optimizer: torch.optim.Optimizer) → None

Log the gradient scaler attribute. If sufficiently high (e.g. >=1024), no gradient clipping required, if it monotonically decaying to 1, add clipping

scaler.get_scale() forces a CUDA sync — only log every 50 steps to avoid serialising the async backward+optimizer pipeline on every step.

on_validation_epoch_end() → None

The following is executed on validation epoch end. Find for the optimal sigmoid activation function threshold that results in the highest Dice score.

test_step(batch: dict[str, torch.Tensor], batch_idx: int) → torch.Tensor

Test step using self.optimal_threshold learned during validation

training_step(batch: dict[str, torch.Tensor], batch_idx: int) → torch.Tensor

Run one training iteration: validate inputs, compute loss, log metrics.

Parameters:

• batch – dict with “image” and “mask” tensors

• batch_idx – index of the current batch within the epoch

Returns:

scalar training loss

Raises:

ValueError – if the batch is missing required keys or contains non-finite values

validation_step(batch: dict[str, torch.Tensor], batch_idx: int) → torch.Tensor

Run one validation iteration and accumulate probabilities and masks for end-of-epoch threshold search.

Parameters:

• batch – dict with “image” and “mask” tensors

• batch_idx – index of the current batch within the epoch

Returns:

scalar validation loss