SemanticSegmentationTask#

class torchgeo.trainers.SemanticSegmentationTask(model='unet', backbone='resnet50', weights=None, in_channels=3, task='multiclass', num_classes=None, num_labels=None, labels=None, num_filters=3, pos_weight=None, loss='ce', class_weights=None, ignore_index=None, lr=0.001, patience=10, freeze_backbone=False, freeze_decoder=False)[source]#

Bases: ClassificationMixin, BaseTask

Semantic Segmentation.

__init__(model='unet', backbone='resnet50', weights=None, in_channels=3, task='multiclass', num_classes=None, num_labels=None, labels=None, num_filters=3, pos_weight=None, loss='ce', class_weights=None, ignore_index=None, lr=0.001, patience=10, freeze_backbone=False, freeze_decoder=False)[source]#

Initialize a new SemanticSegmentationTask instance.

Parameters:

  • model (Literal['unet', 'deeplabv3+', 'fcn', 'upernet', 'segformer', 'dpt']) – Name of the smp model to use.

  • backbone (str) – Name of the timm or smp backbone to use.

  • weights (WeightsEnum | str | bool | None) – Initial model weights. Either a weight enum, the string representation of a weight enum, True for ImageNet weights, False or None for random weights, or the path to a saved model state dict. FCN model does not support pretrained weights.

  • in_channels (int) – Number of input channels to model.

  • task (Literal['binary', 'multiclass', 'multilabel']) – One of ‘binary’, ‘multiclass’, or ‘multilabel’.

  • num_classes (int | None) – Number of prediction classes (only for task='multiclass').

  • num_labels (int | None) – Number of prediction labels (only for task='multilabel').

  • labels (list[str] | None) – List of class names.

  • num_filters (int) – Number of filters. Only applicable when model=’fcn’.

  • pos_weight (Tensor | None) – A weight of positive examples and used with ‘bce’ loss.

  • loss (Literal['ce', 'bce', 'jaccard', 'focal', 'dice']) – Name of the loss function, currently supports ‘ce’, ‘bce’, ‘jaccard’, ‘focal’, and ‘dice’ loss.

  • class_weights (Tensor | Sequence[float] | None) – Optional rescaling weight given to each class and used with ‘ce’ loss.

  • ignore_index (int | None) – Optional integer class index to ignore in the loss and metrics.

  • lr (float) – Learning rate for optimizer.

  • patience (int) – Patience for learning rate scheduler.

  • freeze_backbone (bool) – Freeze the backbone network to fine-tune the decoder and segmentation head.

  • freeze_decoder (bool) – Freeze the decoder network to linear probe the segmentation head.

Added in version 0.9: The labels and pos_weight parameters and dice loss support.

Added in version 0.8: DPT, Segformer, and UPerNet support.

Added in version 0.7: The task and num_labels parameters.

Changed in version 0.6: The ignore_index parameter now works for jaccard loss.

Added in version 0.5: The class_weights, freeze_backbone, and freeze_decoder parameters.

Changed in version 0.5: The weights parameter now supports WeightEnums and checkpoint paths. learning_rate and learning_rate_schedule_patience were renamed to lr and patience.

Changed in version 0.4: segmentation_model, encoder_name, and encoder_weights were renamed to model, backbone, and weights.

Changed in version 0.3: ignore_zeros was renamed to ignore_index.

forward(x)[source]#

Forward pass of the model.

Parameters:

x (Tensor) – Input tensor of shape (B, C, H, W).

Returns:

Output tensor of shape (B, num_classes, H, W).

Return type:

Tensor

configure_models()[source]#

Initialize the model.

training_step(batch, batch_idx, dataloader_idx=0)[source]#

Compute the training loss and additional metrics.

Parameters:

  • batch (dict[str, Any]) – The output of your DataLoader.

  • batch_idx (int) – Integer displaying index of this batch.

  • dataloader_idx (int) – Index of the current dataloader.

Returns:

The loss tensor.

Return type:

Tensor

validation_step(batch, batch_idx, dataloader_idx=0)[source]#

Compute the validation loss and additional metrics.

Parameters:

  • batch (dict[str, Any]) – The output of your DataLoader.

  • batch_idx (int) – Integer displaying index of this batch.

  • dataloader_idx (int) – Index of the current dataloader.

test_step(batch, batch_idx, dataloader_idx=0)[source]#

Compute the test loss and additional metrics.

Parameters:

  • batch (dict[str, Any]) – The output of your DataLoader.

  • batch_idx (int) – Integer displaying index of this batch.

  • dataloader_idx (int) – Index of the current dataloader.

predict_step(batch, batch_idx, dataloader_idx=0)[source]#

Compute the predicted class probabilities.

Parameters:

  • batch (dict[str, Any]) – The output of your DataLoader.

  • batch_idx (int) – Integer displaying index of this batch.

  • dataloader_idx (int) – Index of the current dataloader.

Returns:

Dictionary with ‘probabilities’, ‘bounds’, and ‘transform’ keys.

Return type:

dict[str, Tensor | None]

Changed in version 0.9: Changed return type from Tensor to dict with probabilities, bounds, and transform keys.