MAETask

class torchgeo.trainers.MAETask(model='vit_base_patch32_224', weights=None, in_channels=3, transform=None, decoder_dim=512, lr=0.00015, decoder_num_heads=8, decoder_depth=1, weight_decay=0.05, mask_ratio=0.75, size=224, norm_pix_loss=True, warmup_epochs=40)

Bases: BaseTask

MAE: Masked Autoencoder for self-supervised learning.

Reference implementations:

• facebookresearch/mae

• https://docs.lightly.ai/self-supervised-learning/examples/mae.html

If you use this code for your research, please cite the original paper:

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

ignore = ('transform', 'weights')

Parameters to ignore when saving hyperparameters.

__init__(model='vit_base_patch32_224', weights=None, in_channels=3, transform=None, decoder_dim=512, lr=0.00015, decoder_num_heads=8, decoder_depth=1, weight_decay=0.05, mask_ratio=0.75, size=224, norm_pix_loss=True, warmup_epochs=40)

Initialize the MAE task.

Parameters:

• model (str) – The ViT architecture to use for the encoder. Must be compatible with timm’s create_model function.

• weights (WeightsEnum | str | bool | None) – Pretrained weights to initialize the encoder with. Can be a timm WeightsEnum or a string identifier for a timm weight, True to use default pretrained weights, or None for random initialization.

• in_channels (int) – Number of input channels in the images. Must match the in_chans argument of the ViT model.

• transform (Module | None) – Optional transform to apply to the input images. If None, a default MAE augmentation will be used.

• decoder_dim (int) – The embedding dimension of the MAE decoder. Typically 512 is a good choice for ViT-Base encoders.

• lr (float) – Should typically be set to 1.5e-4 * batch_size / 256.

• decoder_num_heads (int) – Number of attention heads in the MAE decoder.

• decoder_depth (int) – Number of layers in the MAE decoder. Typically 1-4 layers is sufficient for good performance.

• weight_decay (float) – Weight decay for the AdamW optimizer.

• mask_ratio (float) – The ratio of tokens to mask during training. Typically 0.75 is a good choice.

• size (int) – The input image size (height and width) after augmentation. Must match the input size expected by the ViT model.

• norm_pix_loss (bool) – If True, normalize each target patch to zero mean and unit variance before computing MSE. Recommended by the original MAE paper.

• warmup_epochs (int) – Number of linear warmup epochs before cosine annealing.

configure_losses()

Initialize the loss criterion.

configure_models()

Initialize the model.

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(images, idx_keep, idx_mask)

Forward pass through MAE encoder and decoder.

Parameters:

• images (Tensor) – The input images, with shape (B, in_channels, H, W).

• idx_keep (Tensor) – The indices of the tokens that were kept (not masked), with shape (B, N_keep).

• idx_mask (Tensor) – The indices of the tokens that were masked, with shape (B, N_mask).

Returns:

The predicted pixel values for the masked tokens, with shape (B, N_mask,

patch_size*patch_size*in_channels).

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.