SimMIMHead

class mmpretrain.models.heads.SimMIMHead(patch_size, loss)

Head for SimMIM Pre-training.

参数:

• patch_size (int) – Patch size of each token.

• loss (dict) – The config for loss.

loss(pred, target, mask)

Generate loss.

This method will expand mask to the size of the original image.

参数:

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

返回:

The reconstruction loss.

返回类型:

torch.Tensor