MIMHead¶
- class mmpretrain.models.heads.MIMHead(loss)[source]¶
Pre-training head for Masked Image Modeling.
- Parameters:
loss (dict) – Config dict for module of loss functions.
- loss(pred, target, mask=None)[source]¶
Forward head.
- Parameters:
pred (torch.Tensor) – Predictions with shape B x L x C.
target (torch.Tensor) – Targets with shape B x L x C.
mask (torch.Tensor) – Mask with shape B x L.
- Returns:
The loss tensor.
- Return type: