158 lines
6.2 KiB
Python
158 lines
6.2 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from mmengine.dist import all_gather
|
|
from mmengine.model import BaseModule
|
|
|
|
from mmpretrain.registry import MODELS
|
|
|
|
|
|
@MODELS.register_module()
|
|
class ITCHead(BaseModule):
|
|
"""Image-text matching head for multi-modal pre-trained task. Adapted by
|
|
BLIP, ALBEF. Normally used for retrieval task.
|
|
|
|
Args:
|
|
embed_dim (int): Embed channel size for queue.
|
|
queue_size (int): Queue size for image and text. Defaults to 57600.
|
|
temperature (float): Temperature to calculate the similarity.
|
|
Defaults to 0.07.
|
|
use_distill (bool): Whether to use distill to calculate loss.
|
|
Defaults to True.
|
|
alpha (float): Weight for momentum similarity. Defaults to 0.4.
|
|
init_cfg (dict, optional): the config to control the initialization.
|
|
Defaults to None.
|
|
"""
|
|
|
|
def __init__(self,
|
|
embed_dim: int,
|
|
queue_size: int = 57600,
|
|
temperature: float = 0.07,
|
|
use_distill: bool = True,
|
|
alpha: float = 0.4,
|
|
init_cfg: Optional[dict] = None):
|
|
super(ITCHead, self).__init__(init_cfg=init_cfg)
|
|
self.temp = nn.Parameter(temperature * torch.ones([]))
|
|
self.use_distill = use_distill
|
|
if self.use_distill:
|
|
# create the queue
|
|
self.register_buffer('image_queue',
|
|
torch.randn(embed_dim, queue_size))
|
|
self.register_buffer('text_queue',
|
|
torch.randn(embed_dim, queue_size))
|
|
self.register_buffer('idx_queue', torch.full((1, queue_size),
|
|
-100))
|
|
self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))
|
|
|
|
self.image_queue = F.normalize(self.image_queue, dim=0)
|
|
self.text_queue = F.normalize(self.text_queue, dim=0)
|
|
|
|
self.queue_size = queue_size
|
|
# This value will be warmup by `WarmupParamHook`
|
|
self.alpha = alpha
|
|
|
|
def forward(self, feats: Tuple[torch.Tensor]) -> torch.Tensor:
|
|
"""The forward process."""
|
|
return feats[-1]
|
|
|
|
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
|
|
img_feats, text_feats, img_feats_m, text_feats_m = self(feats)
|
|
|
|
img_feats_all = torch.cat(
|
|
[img_feats_m.t(),
|
|
self.image_queue.clone().detach()], dim=1)
|
|
text_feats_all = torch.cat(
|
|
[text_feats_m.t(),
|
|
self.text_queue.clone().detach()], dim=1)
|
|
|
|
# The part can not be traced by torch.fx
|
|
losses = self._get_loss(img_feats, text_feats, img_feats_m,
|
|
text_feats_m, img_feats_all, text_feats_all,
|
|
data_samples, **kwargs)
|
|
return losses
|
|
|
|
def _get_loss(self, img_feats, text_feats, img_feats_m, text_feats_m,
|
|
img_feats_all, text_feats_all, data_samples, **kwargs):
|
|
"""Unpack data samples and compute loss."""
|
|
|
|
idx = torch.tensor([ds.image_id
|
|
for ds in data_samples]).to(img_feats.device)
|
|
idx = idx.view(-1, 1)
|
|
idx_all = torch.cat([idx.t(), self.idx_queue.clone().detach()], dim=1)
|
|
pos_idx = torch.eq(idx, idx_all).float()
|
|
sim_targets = pos_idx / pos_idx.sum(1, keepdim=True)
|
|
|
|
with torch.no_grad():
|
|
if self.use_distill:
|
|
sim_i2t_m = img_feats_m @ text_feats_all / self.temp
|
|
sim_t2i_m = text_feats_m @ img_feats_all / self.temp
|
|
|
|
sim_i2t_targets = (
|
|
self.alpha * F.softmax(sim_i2t_m, dim=1) +
|
|
(1 - self.alpha) * sim_targets)
|
|
sim_t2i_targets = (
|
|
self.alpha * F.softmax(sim_t2i_m, dim=1) +
|
|
(1 - self.alpha) * sim_targets)
|
|
|
|
sim_i2t = img_feats @ text_feats_all / self.temp
|
|
sim_t2i = text_feats @ img_feats_all / self.temp
|
|
|
|
if self.use_distill:
|
|
loss_i2t = -torch.sum(
|
|
F.log_softmax(sim_i2t, dim=1) * sim_i2t_targets, dim=1).mean()
|
|
loss_t2i = -torch.sum(
|
|
F.log_softmax(sim_t2i, dim=1) * sim_t2i_targets, dim=1).mean()
|
|
else:
|
|
loss_i2t = -torch.sum(
|
|
F.log_softmax(sim_i2t, dim=1) * sim_targets, dim=1).mean()
|
|
loss_t2i = -torch.sum(
|
|
F.log_softmax(sim_t2i, dim=1) * sim_targets, dim=1).mean()
|
|
|
|
# compute loss
|
|
losses = dict()
|
|
|
|
losses['itc_loss'] = (loss_i2t + loss_t2i) / 2
|
|
self._dequeue_and_enqueue(img_feats_m, text_feats_m, idx)
|
|
return losses
|
|
|
|
@torch.no_grad()
|
|
def _dequeue_and_enqueue(self, image_feat, text_feat, idxs=None):
|
|
# gather keys before updating queue
|
|
image_feats = torch.cat(all_gather(image_feat))
|
|
text_feats = torch.cat(all_gather(text_feat))
|
|
|
|
batch_size = image_feats.shape[0]
|
|
|
|
ptr = int(self.queue_ptr)
|
|
assert self.queue_size % batch_size == 0 # for simplicity
|
|
|
|
# replace the keys at ptr (dequeue and enqueue)
|
|
self.image_queue[:, ptr:ptr + batch_size] = image_feats.T
|
|
self.text_queue[:, ptr:ptr + batch_size] = text_feats.T
|
|
|
|
if idxs is not None:
|
|
idxs = torch.cat(all_gather(idxs))
|
|
self.idx_queue[:, ptr:ptr + batch_size] = idxs.T
|
|
|
|
ptr = (ptr + batch_size) % self.queue_size # move pointer
|
|
self.queue_ptr[0] = ptr
|