SpatioTemporalSegmentationTask

class torchgeo.trainers.SpatioTemporalSegmentationTask(model='convlstm', in_channels=3, task='multiclass', num_classes=None, num_labels=None, labels=None, pos_weight=None, loss='ce', class_weights=None, ignore_index=None, lr=0.001, patience=10, **kwargs)

Bases: ClassificationMixin, BaseTask

Spatiotemporal Semantic Segmentation.

Added in version 0.10.

__init__(model='convlstm', in_channels=3, task='multiclass', num_classes=None, num_labels=None, labels=None, pos_weight=None, loss='ce', class_weights=None, ignore_index=None, lr=0.001, patience=10, **kwargs)

Initialize a new SpatioTemporalSegmentationTask instance.

Parameters:

• model (Literal['convlstm']) – Spatiotemporal model name. Supported value is 'convlstm'.

• in_channels (int) – Number of channels per timestep for inputs of shape (B, T, C, H, W).

• 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.

• 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.

• **kwargs (Any) – Additional keyword arguments passed to the model constructor.

forward(x, **kwargs)

Forward pass of the model.

Parameters:

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

• **kwargs (Any) – Additional keyword arguments forwarded to the model.

Returns:

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

Return type:

Tensor

configure_models()

Initialize the model.

training_step(batch, batch_idx, dataloader_idx=0)

Compute the training loss and additional metrics.

validation_step(batch, batch_idx, dataloader_idx=0)

Compute the validation loss and additional metrics.

test_step(batch, batch_idx, dataloader_idx=0)

Compute the test loss and additional metrics.

predict_step(batch, batch_idx, dataloader_idx=0)

Compute the predicted class probabilities.