RegressionTask#

class torchgeo.trainers.RegressionTask(model='resnet50', backbone='resnet50', weights=None, in_channels=3, num_outputs=1, labels=None, num_filters=3, loss='mse', lr=0.001, patience=10, freeze_backbone=False, freeze_decoder=False)[source]#

Bases: RegressionMixin, BaseTask

Regression.

__init__(model='resnet50', backbone='resnet50', weights=None, in_channels=3, num_outputs=1, labels=None, num_filters=3, loss='mse', lr=0.001, patience=10, freeze_backbone=False, freeze_decoder=False)[source]#

Initialize a new RegressionTask instance.

Parameters:

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

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

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

  • num_outputs (int) – Number of prediction outputs.

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

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

  • loss (Literal['mae', 'mse']) – One of ‘mse’ or ‘mae’.

  • lr (float) – Learning rate for optimizer.

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

  • freeze_backbone (bool) – Freeze the backbone network to linear probe the regression head. Does not support FCN models.

  • freeze_decoder (bool) – Freeze the decoder network to linear probe the regression head. Does not support FCN models. Only applicable to PixelwiseRegressionTask.

Changed in version 0.4: Change regression model support from torchvision.models to timm

Added in version 0.5: The freeze_backbone and freeze_decoder parameters.

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

Added in version 0.10: The labels parameter.

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 regression values.

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