BYOLTask

class torchgeo.trainers.BYOLTask(model='resnet50', weights=None, in_channels=3, lr=0.001, patience=10)

Bases: BaseTask

BYOL: Bootstrap Your Own Latent.

Reference implementation:

• google-deepmind/deepmind-research

If you use this trainer in your research, please cite the following paper:

• https://arxiv.org/abs/2006.07733

monitor = 'train_loss'

Performance metric to monitor in learning rate scheduler and callbacks.

__init__(model='resnet50', weights=None, in_channels=3, lr=0.001, patience=10)

Initialize a new BYOLTask 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.

• lr (float) – Learning rate for optimizer.

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

Changed in version 0.4: backbone_name was renamed to backbone. Changed backbone support from torchvision.models to timm.

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

configure_models()

Initialize the model.

training_step(batch, batch_idx, dataloader_idx=0)

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.

Raises:

AssertionError – If channel dimensions are incorrect.

Return type:

Tensor

validation_step(batch, batch_idx, dataloader_idx=0)

No-op, does nothing.

test_step(batch, batch_idx, dataloader_idx=0)

No-op, does nothing.

predict_step(batch, batch_idx, dataloader_idx=0)

No-op, does nothing.