MixMIMPretrainHead

class mmpretrain.models.heads.MixMIMPretrainHead(loss, norm_pix=False, patch_size=16)

Head for MixMIM Pre-training.

Parameters:

• loss (dict) – Config of loss.

• norm_pix_loss (bool) – Whether or not normalize target. Defaults to False.

• patch_size (int) – Patch size. Defaults to 16.

loss(x_rec, target, mask)

Generate loss.

Parameters:

• pred (torch.Tensor) – The reconstructed image.

• target (torch.Tensor) – The target image.

• mask (torch.Tensor) – The mask of the target image.

Returns:

The reconstruction loss.

Return type:

torch.Tensor