MoCoTask

class torchgeo.trainers.MoCoTask(model='resnet50', weights=None, in_channels=3, version=3, layers=3, hidden_dim=4096, output_dim=256, lr=9.6, weight_decay=1e-06, momentum=0.9, schedule=[120, 160], temperature=1, memory_bank_size=0, moco_momentum=0.99, gather_distributed=False, size=224, grayscale_weights=None, augmentation1=None, augmentation2=None)

Bases: BaseTask

MoCo: Momentum Contrast.

Reference implementations:

• facebookresearch/moco

• facebookresearch/moco-v3

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

• v1: https://arxiv.org/abs/1911.05722

• v2: https://arxiv.org/abs/2003.04297

• v3: https://arxiv.org/abs/2104.02057

Added in version 0.5.

ignore = ('weights', 'augmentation1', 'augmentation2')

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=3, layers=3, hidden_dim=4096, output_dim=256, lr=9.6, weight_decay=1e-06, momentum=0.9, schedule=[120, 160], temperature=1, memory_bank_size=0, moco_momentum=0.99, gather_distributed=False, size=224, grayscale_weights=None, augmentation1=None, augmentation2=None)

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

• version (int) – Version of MoCo, 1–3.

• layers (int) – Number of layers in projection head (not used in v1, 2 for v1/2, 3 for v3).

• hidden_dim (int) – Number of hidden dimensions in projection head (not used in v1, 2048 for v2, 4096 for v3).

• output_dim (int) – Number of output dimensions in projection head (not used in v1, 128 for v2, 256 for v3).

• lr (float) – Learning rate (0.03 x batch_size / 256 for v1/2, 0.6 x batch_size / 256 for v3).

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

• momentum (float) – Momentum of SGD solver (v1/2 only).

• schedule (Sequence[int]) – Epochs at which to drop lr by 10x (v1/2 only).

• temperature (float) – Temperature used in InfoNCE loss (0.07 for v1/2, 1 for v3).

• memory_bank_size (int) – Size of memory bank (65536 for v1/2, 0 for v3).

• moco_momentum (float) – MoCo momentum of updating key encoder (0.999 for v1/2, 0.99 for v3)

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

• augmentation1 (Module | None) – Data augmentation for 1st branch. Defaults to MoCo augmentation.

• augmentation2 (Module | None) – Data augmentation for 2nd branch. Defaults to MoCo augmentation.

Raises:

AssertionError – If an invalid version of MoCo is requested.

Warns:

UserWarning – If hyperparameters do not match MoCo version requested.

configure_models()

Initialize the model.

configure_losses()

Initialize the loss criterion.

configure_optimizers()

Initialize the optimizer and learning rate scheduler.

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

forward(x)

Forward pass of the model.

Parameters:

x (Tensor) – Mini-batch of images.

Returns:

Output of the model and backbone

Return type:

tuple[Tensor, Tensor]

forward_momentum(x)

Forward pass of the momentum model.

Parameters:

x (Tensor) – Mini-batch of images.

Returns:

Output from the momentum model.

Return type:

Tensor

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.

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.