ObjectDetectionTask#

class torchgeo.trainers.ObjectDetectionTask(model='faster-rcnn', backbone='resnet50', weights=None, in_channels=3, num_classes=1000, trainable_layers=3, lr=0.001, patience=10, freeze_backbone=False)[source]#

Bases: BaseTask

Object detection.

Added in version 0.4.

monitor = 'val_map'#

Performance metric to monitor in learning rate scheduler and callbacks.

mode = 'max'#

Whether the goal is to minimize or maximize the performance metric to monitor.

__init__(model='faster-rcnn', backbone='resnet50', weights=None, in_channels=3, num_classes=1000, trainable_layers=3, lr=0.001, patience=10, freeze_backbone=False)[source]#

Initialize a new ObjectDetectionTask instance.

Note that we disable the internal normalize+resize transform of the detection models. Please ensure your images are appropriately resized before passing them to the model.

Parameters:
  • model (str) – Name of the torchvision model to use. One of ‘faster-rcnn’, ‘fcos’, or ‘retinanet’.

  • backbone (str) – Name of the torchvision backbone to use. One of ‘resnet18’, ‘resnet34’, ‘resnet50’, ‘resnet101’, ‘resnet152’, ‘resnext50_32x4d’, ‘resnext101_32x8d’, ‘wide_resnet50_2’, or ‘wide_resnet101_2’.

  • weights (WeightsEnum | None) – Initial model weights.

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

  • num_classes (int) – Number of prediction classes (including the background).

  • trainable_layers (int) – Number of trainable layers.

  • lr (float) – Learning rate for optimizer.

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

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

Changed in version 0.4: detection_model was renamed to model.

Added in version 0.5: The freeze_backbone parameter.

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

configure_models()[source]#

Initialize the model.

Raises:

ValueError – If model or backbone are invalid.

configure_metrics()[source]#

Initialize the performance metrics.

  • MeanAveragePrecision: Mean average precision (mAP) and mean average recall (mAR). Precision is the number of true positives divided by the number of true positives + false positives. Recall is the number of true positives divived by the number of true positives + false negatives. Uses ‘macro’ averaging. Higher values are better.

Note

  • ‘Micro’ averaging suits overall performance evaluation but may not reflect minority class accuracy.

  • ‘Macro’ averaging gives equal weight to each class, and is useful for balanced performance assessment across imbalanced classes.

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

Compute the training loss.

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 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 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 bounding boxes.

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:

list[dict[str, Tensor]]