MIMHead

class mmpretrain.models.heads.MIMHead(loss)

Pre-training head for Masked Image Modeling.

参数:

loss (dict) – Config dict for module of loss functions.

loss(pred, target, mask=None)

Forward head.

参数:

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

返回:

The loss tensor.

返回类型:

torch.Tensor