SimMIMHead¶
- class mmpretrain.models.heads.SimMIMHead(patch_size, loss)[source]¶
Head for SimMIM Pre-training.
- loss(pred, target, mask)[source]¶
Generate loss.
This method will expand mask to the size of the original image.
- Parameters:
pred (torch.Tensor) – The reconstructed image (B, C, H, W).
target (torch.Tensor) – The target image (B, C, H, W).
mask (torch.Tensor) – The mask of the target image.
- Returns:
The reconstruction loss.
- Return type: