Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from typing import Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| from mmengine.model import BaseModule | |
| from mmpretrain.evaluation import Accuracy | |
| from mmpretrain.registry import MODELS | |
| class Pooler(nn.Module): | |
| def __init__(self, hidden_size): | |
| super().__init__() | |
| self.dense = nn.Linear(hidden_size, hidden_size) | |
| self.activation = nn.Tanh() | |
| def forward(self, hidden_states): | |
| first_token_tensor = hidden_states[:, 0] | |
| pooled_output = self.dense(first_token_tensor) | |
| pooled_output = self.activation(pooled_output) | |
| return pooled_output | |
| class ITMHead(BaseModule): | |
| """Image-text matching head for multi-modal pre-trained task. Adapted by | |
| BLIP, FLAVA. | |
| Args: | |
| hidden_size (int): Hidden channel size out input features. | |
| with_pooler (bool): Whether a pooler is added. Defaults to True. | |
| loss (dict): Config of global contrasive loss. Defaults to | |
| ``dict(type='GlobalContrasiveLoss')``. | |
| cal_acc (bool): Whether to calculate accuracy during training. | |
| If you use batch augmentations like Mixup and CutMix during | |
| training, it is pointless to calculate accuracy. | |
| Defaults to False. | |
| init_cfg (dict, optional): the config to control the initialization. | |
| Defaults to None. | |
| """ | |
| def __init__(self, | |
| hidden_size: int, | |
| with_pooler: bool = True, | |
| loss: dict = dict(type='CrossEntropyLoss', loss_weight=1.0), | |
| cal_acc: bool = False, | |
| init_cfg: Optional[dict] = None): | |
| super(ITMHead, self).__init__(init_cfg=init_cfg) | |
| self.hidden_size = hidden_size | |
| if with_pooler: | |
| self.pooler = Pooler(hidden_size=self.hidden_size) | |
| else: | |
| self.pooler = nn.Identity() | |
| self.fc = nn.Linear(self.hidden_size, 2) | |
| self.loss_module = MODELS.build(loss) | |
| self.cal_acc = cal_acc | |
| def forward(self, feats: Tuple[torch.Tensor]) -> torch.Tensor: | |
| """The forward process.""" | |
| pre_logits = self.pooler(feats[-1]) | |
| itm_logits = self.fc(pre_logits) | |
| return itm_logits | |
| def loss(self, feats: Tuple[torch.Tensor], data_samples, **kwargs) -> dict: | |
| """Calculate losses from the classification score. | |
| Args: | |
| feats (tuple[Tensor]): The features extracted from the backbone. | |
| Multiple stage inputs are acceptable but only the last stage | |
| will be used to classify. The shape of every item should be | |
| ``(num_samples, num_classes)``. | |
| data_samples (List[ClsDataSample]): The annotation data of | |
| every samples. | |
| **kwargs: Other keyword arguments to forward the loss module. | |
| Returns: | |
| dict[str, Tensor]: a dictionary of loss components | |
| """ | |
| # The part can be traced by torch.fx | |
| itm_logits = self(feats) | |
| # deal with query | |
| if itm_logits.ndim == 3: | |
| itm_logits = itm_logits.mean(dim=1) | |
| # The part can not be traced by torch.fx | |
| losses = self._get_loss(itm_logits, data_samples, **kwargs) | |
| return losses | |
| def _get_loss(self, itm_logits: torch.Tensor, data_samples, **kwargs): | |
| """Unpack data samples and compute loss.""" | |
| # Unpack data samples and pack targets | |
| # use `itm_label` in here temporarily | |
| target = torch.tensor([i.is_matched | |
| for i in data_samples]).to(itm_logits.device) | |
| # compute loss | |
| losses = dict() | |
| loss = self.loss_module( | |
| itm_logits, target.long(), avg_factor=itm_logits.size(0), **kwargs) | |
| losses['itm_loss'] = loss | |
| # compute accuracy | |
| if self.cal_acc: | |
| # topk is meaningless for matching task | |
| acc = Accuracy.calculate(itm_logits, target) | |
| # acc is warpped with two lists of topk and thrs | |
| # which are unnecessary here | |
| losses.update({'itm_accuracy': acc[0][0]}) | |
| return losses | |