SegformerFinetuner

The lightning_model module provides a PyTorch Lightning LightningModule for fine-tuning the Segformer model for semantic segmentation tasks.

class ucs.model.lightning_model.SegformerFinetuner(config=None, class_weights=None, **kwargs)

Bases: LightningModule

A PyTorch Lightning module for fine-tuning the SegFormer model for semantic segmentation tasks.

Variables:

• model (SegformerForSemanticSegmentation) – The SegFormer model for semantic segmentation.

• metrics (SegMetrics) – Metrics object for tracking performance.

• criterion (CeDiceLoss) – The loss function used for training.

• test_results (dict) – Stores ‘predictions’ and ‘ground_truths’ for test evaluation.

Parameters:

• config (TrainingConfig, optional) – Training configuration containing model hyperparameters, stored in self.hparams.

• class_weights (Tensor, optional) – Class weights for loss balancing.

• **kwargs (dict) –

  Additional hyperparameters that override config, such as:

  • model_name (str): The SegFormer variant to use (e.g., “b0”).

  • max_epochs (int): Maximum number of training epochs.

  • learning_rate (float): Learning rate for the optimizer.

  • weight_decay (float): Weight decay (L2 regularization).

  • ignore_index (Optional[int]): Label index to ignore during training.

  • weighting_strategy (str): Strategy for class weighting (‘none’, ‘balanced’, ‘max’, ‘sum’, or ‘raw’).

  • alpha (float): CeDiceLoss loss alpha parameter.

  • id2label (Dict[int, str]): Mapping from class indices to class labels.

calculate_confusion_matrix()

Calculate the confusion matrix from test predictions and ground truths.

Returns:

The confusion matrix.

Return type:

ndarray

configure_optimizers()

Set up the optimizer and learning rate scheduler.

Returns:

A list containing:

• optimizer (torch.optim.AdamW): Optimizer configured with weight decay and learning rate.

• scheduler (dict): CosineAnnealingLR scheduler, which reduces the learning rate following a cosine schedule.

Return type:

tuple

forward(images, masks=None)

Forward pass through the model.

Parameters:

• images (Tensor) – Input images of shape (batch_size, num_channels, height, width).

• masks (Tensor, optional) – Ground truth masks of shape (batch_size, height, width). Default is None.

Returns:

loss (Tensor or None): Computed loss if masks are provided, else None. predictions (Tensor): Predicted labels of shape (batch_size, height, width), with class indices for each pixel.

Return type:

tuple

freeze_encoder_layers(blocks_to_freeze=None)

Freezes specified encoder layers to prevent weight updates.

Parameters:

blocks_to_freeze (list`[:class:`str], optional) – List of encoder blocks to freeze.

log_step(stage, loss, metrics)

Logs loss and evaluation metrics.

Parameters:

• stage (str) – One of ‘train’, ‘val’, or ‘test’.

• loss (Tensor) – Loss value for the current step.

• metrics (dict) – Computed segmentation metrics.

on_fit_start()

Sets the model in training mode at the start of training.

on_test_epoch_end()

Reset the metrics at the end of the test epoch.

on_test_start()

Add test-specific metrics at the start of the test phase.

on_train_epoch_end()

Reset the metrics at the end of the train epoch.

This prevents the accumulation of metric values across epochs and ensures metrics are calculated independently for each epoch.

on_train_start()

Called at the start of training, set model in training mode.

on_validation_epoch_end()

Reset the metrics at the end of the validation epoch.

This prevents the accumulation of metric values across epochs and ensures metrics are calculated independently for each epoch.

save_pretrained_model(pretrained_path)

Save the trained model in Hugging Face’s Transformers-compatible format.

Parameters:

pretrained_path (str or Path) – Directory where the model will be saved.

Notes

This allows the saved model to be loaded later using the from_pretrained method.

step(batch, stage)

Runs a single step of training, validation, or testing.

Parameters:

• batch (dict) – Contains ‘pixel_values’ (images) and ‘labels’ (masks).

• stage (str) – One of ‘train’, ‘val’, or ‘test’.

Returns:

The computed loss.

Return type:

Tensor

test_step(batch, batch_idx)

Perform a single test step.

Parameters:

• batch (dict) – Contains input images and ground truth masks.

• batch_idx (int) – Index of the current batch.

Returns:

Loss, predictions, and ground truths for further aggregation.

Return type:

dict

training_step(batch, batch_idx)

Perform a single training step.

Parameters:

• batch (dict) – A batch of data containing ‘pixel_values’ and ‘labels’.

• batch_idx (int) – The index of the batch.

Returns:

The computed loss.

Return type:

Tensor

unfreeze_encoder_layers(blocks_to_unfreeze=None)

Unfreezes specified encoder layers to allow training updates.

Parameters:

blocks_to_unfreeze (list`[:class:`str], optional) – List of encoder blocks to unfreeze.

validation_step(batch, batch_idx)

Perform a single validation step.

Parameters:

• batch (dict) – A batch of data containing ‘pixel_values’ and ‘labels’.

• batch_idx (int) – The index of the batch.

Returns:

The computed loss.

Return type:

Tensor