SimCLRTask#

class torchgeo.trainers.SimCLRTask(model='resnet50', weights=None, in_channels=3, version=2, layers=3, hidden_dim=None, output_dim=None, lr=4.8, momentum=0.9, weight_decay=0.0001, temperature=0.07, memory_bank_size=64000, gather_distributed=False, size=224, grayscale_weights=None, augmentations=None)[source]#

Bases: BaseTask

SimCLR: a simple framework for contrastive learning of visual representations.

Reference implementation:

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

Added in version 0.5.

ignore = ('weights', 'augmentations')#

Parameters to ignore when saving hyperparameters.

monitor = 'train_loss'#

Performance metric to monitor in learning rate scheduler and callbacks.

__init__(model='resnet50', weights=None, in_channels=3, version=2, layers=3, hidden_dim=None, output_dim=None, lr=4.8, momentum=0.9, weight_decay=0.0001, temperature=0.07, memory_bank_size=64000, gather_distributed=False, size=224, grayscale_weights=None, augmentations=None)[source]#

Initialize a new SimCLRTask instance.

Added in version 0.6: The momentum parameter.

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.

  • version (int) – Version of SimCLR, 1–2.

  • layers (int) – Number of layers in projection head (2 for v1, 3+ for v2).

  • hidden_dim (int | None) – Number of hidden dimensions in projection head (defaults to output dimension of model).

  • output_dim (int | None) – Number of output dimensions in projection head (defaults to output dimension of model).

  • lr (float) – Learning rate (0.3 x batch_size / 256 is recommended).

  • momentum (float) – Momentum factor.

  • weight_decay (float) – Weight decay coefficient (1e-6 for v1, 1e-4 for v2).

  • temperature (float) – Temperature used in NT-Xent loss.

  • memory_bank_size (int) – Size of memory bank (0 for v1, 64K for v2).

  • gather_distributed (bool) – Gather negatives from all GPUs during distributed training (ignored if memory_bank_size > 0).

  • size (int) – Size of patch to crop.

  • grayscale_weights (Tensor | None) – Weight vector for grayscale computation, see RandomGrayscale. Only used when augmentations=None. Defaults to average of all bands.

  • augmentations (Module | None) – Data augmentation. Defaults to SimCLR augmentation.

Raises:

AssertionError – If an invalid version of SimCLR is requested.

Warns:

UserWarning – If hyperparameters do not match SimCLR version requested.

configure_models()[source]#

Initialize the model.

configure_losses()[source]#

Initialize the loss criterion.

forward(x)[source]#

Forward pass of the model.

Parameters:

x (Tensor) – Mini-batch of images.

Returns:

Output of the model and backbone.

Return type:

tuple[Tensor, Tensor]

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.

Raises:

AssertionError – If channel dimensions are incorrect.

Return type:

Tensor

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

No-op, does nothing.

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

No-op, does nothing.

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

No-op, does nothing.

configure_optimizers()[source]#

Initialize the optimizer and learning rate scheduler.

Changed in version 0.6: Changed from Adam to LARS optimizer.

Returns:

Optimizer and learning rate scheduler.

Return type:

Optimizer | Sequence[Optimizer] | tuple[Sequence[Optimizer], Sequence[LRScheduler | ReduceLROnPlateau | LRSchedulerConfig]] | OptimizerConfig | OptimizerLRSchedulerConfig | Sequence[OptimizerConfig] | Sequence[OptimizerLRSchedulerConfig] | None