ClassificationTask#

class torchgeo.trainers.ClassificationTask(model='resnet50', weights=None, 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, freeze_backbone=False)[source]#

Bases: ClassificationMixin, BaseTask

Image classification.

__init__(model='resnet50', weights=None, 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, freeze_backbone=False)[source]#

Initialize a new ClassificationTask instance.

Parameters:

  • model (str) – Name of the timm model 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.

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

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

  • loss (Literal['ce', 'bce', 'jaccard', 'focal', 'dice']) – One of ‘ce’, ‘bce’, ‘jaccard’, ‘focal’, or ‘dice’.

  • 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 linear probe the classifier head.

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

Added in version 0.7: The task and num_labels parameters.

Added in version 0.5: The class_weights and freeze_backbone parameters.

Changed in version 0.5: learning_rate and learning_rate_schedule_patience were renamed to lr and patience.

Changed in version 0.4: classification_model was renamed to model.

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:

Output predicted probabilities.

Return type:

Tensor