38 lines
1003 B
Python
38 lines
1003 B
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from mmengine.model import BaseModule
|
|
|
|
from mmpretrain.registry import MODELS
|
|
|
|
|
|
@MODELS.register_module()
|
|
class MIMHead(BaseModule):
|
|
"""Pre-training head for Masked Image Modeling.
|
|
|
|
Args:
|
|
loss (dict): Config dict for module of loss functions.
|
|
"""
|
|
|
|
def __init__(self, loss: dict) -> None:
|
|
super().__init__()
|
|
self.loss_module = MODELS.build(loss)
|
|
|
|
def loss(self,
|
|
pred: torch.Tensor,
|
|
target: torch.Tensor,
|
|
mask: Optional[torch.Tensor] = None) -> torch.Tensor:
|
|
"""Forward head.
|
|
|
|
Args:
|
|
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:
|
|
torch.Tensor: The loss tensor.
|
|
"""
|
|
loss = self.loss_module(pred, target, mask)
|
|
return loss
|