506 lines
21 KiB
Python
506 lines
21 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from typing import Dict, List, Optional, Tuple, Union
|
|
|
|
import mmengine.dist as dist
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from mmengine.utils import track_iter_progress
|
|
|
|
from mmpretrain.registry import MODELS, TOKENIZER
|
|
from mmpretrain.structures import DataSample
|
|
from ..blip.blip_retrieval import BlipRetrieval, all_gather_concat
|
|
|
|
|
|
@MODELS.register_module()
|
|
class Blip2Retrieval(BlipRetrieval):
|
|
"""BLIP2 Retriever.
|
|
|
|
Args:
|
|
vision_backbone (dict): Backbone for extracting image features.
|
|
text_backbone (dict): Backbone for extracting text features.
|
|
multimodal_backbone (Optional[dict]): Backbone for extracting
|
|
multi-modal features.
|
|
vision_neck (Optional[dict]): The neck module to process image features
|
|
from vision backbone. Defaults to None.
|
|
text_neck (Optional[dict]): The neck module to process text features
|
|
from text backbone. Defaults to None.
|
|
head (Optional[Union[List[dict], dict]]): The head module to calculate
|
|
loss from processed single modality features.
|
|
See :mod:`mmmultimodal.models.heads`.
|
|
Notice that if the head is not set, `loss` method cannot be used.
|
|
Defaults to None.
|
|
multimodal_head (Optional[Union[List[dict], dict]]): The multi-modal
|
|
head module to calculate loss from processed multimodal features.
|
|
See :mod:`mmmultimodal.models.heads`.
|
|
Notice that if the head is not set, `loss` method cannot be used.
|
|
Defaults to None.
|
|
tokenizer (Optional[dict]): The config for tokenizer. Defaults to None.
|
|
temperature (float): Temperature parameter that controls the
|
|
concentration level of the distribution. Defaults to 0.07.
|
|
fast_match (bool): If False, select topk similarity as candidates and
|
|
compute the matching score. If True, return the similarity as the
|
|
matching score directly. Defaults to False.
|
|
topk (int): Select topk similarity as candidates for compute matching
|
|
scores. Notice that this is not the topk in evaluation.
|
|
Defaults to 256.
|
|
data_preprocessor (Optional[dict]): The config for preprocessing input
|
|
data. If None or no specified type, it will use
|
|
"MultiModalDataPreprocessor" as type.
|
|
See :class:`MultiModalDataPreprocessor` for more details.
|
|
Defaults to None.
|
|
init_cfg (Optional[dict]): the config to control the initialization.
|
|
Defaults to None.
|
|
"""
|
|
|
|
def __init__(self,
|
|
vision_backbone: dict,
|
|
text_backbone: Optional[dict] = None,
|
|
multimodal_backbone: Optional[dict] = None,
|
|
vision_neck: Optional[dict] = None,
|
|
text_neck: Optional[dict] = None,
|
|
head: Optional[Union[List[dict], dict]] = None,
|
|
multimodal_head: Optional[Union[List[dict], dict]] = None,
|
|
tokenizer: Optional[dict] = None,
|
|
temperature: float = 0.07,
|
|
fast_match: bool = False,
|
|
topk: int = 256,
|
|
data_preprocessor: Optional[dict] = None,
|
|
init_cfg: Optional[dict] = None) -> None:
|
|
if data_preprocessor is None:
|
|
data_preprocessor = {}
|
|
if isinstance(data_preprocessor, dict):
|
|
data_preprocessor.setdefault('type', 'MultiModalDataPreprocessor')
|
|
data_preprocessor = MODELS.build(data_preprocessor)
|
|
|
|
# Skip BlipRetrieval init
|
|
super(BlipRetrieval, self).__init__(
|
|
init_cfg=init_cfg, data_preprocessor=data_preprocessor)
|
|
|
|
self.vision_backbone = MODELS.build(vision_backbone)
|
|
self.ln_vision_backbone = nn.LayerNorm(self.vision_backbone.embed_dims)
|
|
self.tokenizer = TOKENIZER.build(tokenizer)
|
|
|
|
if text_backbone is not None:
|
|
self.text_backbone = MODELS.build(text_backbone)
|
|
|
|
if multimodal_backbone is not None:
|
|
self.multimodal_backbone = MODELS.build(multimodal_backbone)
|
|
self.multimodal_backbone.resize_token_embeddings(
|
|
len(self.tokenizer))
|
|
self.query_tokens = nn.Parameter(
|
|
torch.zeros(1, self.multimodal_backbone.bert.config.query_length,
|
|
self.multimodal_backbone.bert.config.hidden_size))
|
|
self.query_tokens.data.normal_(
|
|
mean=0.0,
|
|
std=self.multimodal_backbone.bert.config.initializer_range)
|
|
|
|
if vision_neck is not None:
|
|
self.vision_neck = MODELS.build(vision_neck)
|
|
|
|
if text_neck is not None:
|
|
self.text_neck = MODELS.build(text_neck)
|
|
|
|
if head is not None:
|
|
self.head = MODELS.build(head)
|
|
|
|
if multimodal_head is not None:
|
|
self.multimodal_head = MODELS.build(multimodal_head)
|
|
|
|
self.temp = nn.Parameter(temperature * torch.ones([]))
|
|
|
|
# Notice that this topk is used for select k candidate to compute
|
|
# image-text score, but not the final metric topk in evaluation.
|
|
self.fast_match = fast_match
|
|
self.topk = topk
|
|
|
|
def _extract_feat(self, inputs: Union[torch.Tensor, dict],
|
|
modality: str) -> Tuple[torch.Tensor]:
|
|
"""Extract features from the single modality.
|
|
Args:
|
|
inputs (Union[torch.Tensor, dict]): A batch of inputs.
|
|
For image, a tensor of shape (N, C, ...) in general.
|
|
For text, a dict of tokenized text inputs.
|
|
modality (str): Modality feature to be extracted. Only two
|
|
options are supported.
|
|
|
|
- ``images``: Only extract image features, mostly used for
|
|
inference.
|
|
- ``texts``: Only extract text features, mostly used for
|
|
inference.
|
|
Returns:
|
|
Tuple[torch.Tensor]: The output features.
|
|
"""
|
|
if modality == 'images':
|
|
# extract image features
|
|
# TODO:
|
|
# Add layernorm inside backbone and handle the concat outside
|
|
image_embeds = self.ln_vision_backbone(
|
|
self.vision_backbone(inputs)[0])
|
|
image_atts = torch.ones(
|
|
image_embeds.size()[:-1], dtype=torch.long).to(self.device)
|
|
|
|
query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1,
|
|
-1)
|
|
query_output = self.multimodal_backbone.bert(
|
|
query_embeds=query_tokens,
|
|
encoder_hidden_states=image_embeds,
|
|
encoder_attention_mask=image_atts,
|
|
use_cache=True,
|
|
return_dict=True,
|
|
)
|
|
image_feat = F.normalize(
|
|
self.vision_neck([query_output.last_hidden_state]), dim=-1)
|
|
return {
|
|
'image_embeds': image_embeds,
|
|
'image_feat': image_feat,
|
|
'query_output': query_output
|
|
}
|
|
elif modality == 'texts':
|
|
# extract text features
|
|
text_output = self.multimodal_backbone.bert(
|
|
inputs.input_ids,
|
|
attention_mask=inputs.attention_mask,
|
|
return_dict=True,
|
|
)
|
|
text_embeds = text_output.last_hidden_state
|
|
text_feat = F.normalize(
|
|
self.text_neck([text_embeds[:, 0, :]]), dim=-1)
|
|
return {'text_embeds': text_embeds, 'text_feat': text_feat}
|
|
else:
|
|
raise RuntimeError(f'Invalid modality "{modality}".')
|
|
|
|
def loss(
|
|
self,
|
|
images: torch.Tensor,
|
|
data_samples: Optional[List[DataSample]] = None,
|
|
) -> Dict[str, torch.tensor]:
|
|
"""Calculate losses from a batch of inputs and data samples.
|
|
|
|
Args:
|
|
inputs (dict): A batch of inputs. The input tensor with of
|
|
at least one modality. For image, the value is a tensor
|
|
of shape (N, C, ...) in general.
|
|
For text, the value is a dict of tokenized text inputs.
|
|
data_samples (Optional[List[DataSample]]):
|
|
The annotation data of every samples. Defaults to None.
|
|
|
|
Returns:
|
|
Dict[str, torch.tensor]: a dictionary of loss components of
|
|
both head and multimodal head.
|
|
"""
|
|
output = self.extract_feat(images, data_samples)
|
|
|
|
text_ids = output['text_ids']
|
|
text_attn_mask = output['text_attn_mask']
|
|
image_embeds = output['image_embeds']
|
|
image_feat = output['image_feat']
|
|
text_feat = output['text_feat']
|
|
query_output = output['query_output']
|
|
|
|
# ITC Loss
|
|
# B*world_size, num_query, D
|
|
image_feat_all = torch.cat(dist.all_gather(image_feat))
|
|
# B*world_size, D
|
|
text_feat_all = torch.cat(dist.all_gather(text_feat))
|
|
|
|
# B, B*world_size, num_query
|
|
sim_q2t = torch.matmul(
|
|
image_feat.unsqueeze(1), text_feat_all.unsqueeze(-1)).squeeze()
|
|
|
|
# image to text similarity
|
|
sim_i2t, _ = sim_q2t.max(-1)
|
|
sim_i2t = sim_i2t / self.temp
|
|
|
|
# B, B*world_size, num_query
|
|
sim_t2q = torch.matmul(
|
|
text_feat.unsqueeze(1).unsqueeze(1),
|
|
image_feat_all.permute(0, 2, 1)).squeeze()
|
|
|
|
# text-image similarity
|
|
sim_t2i, _ = sim_t2q.max(-1)
|
|
sim_t2i = sim_t2i / self.temp
|
|
|
|
rank = dist.get_rank()
|
|
bs = images.size(0)
|
|
targets = torch.linspace(
|
|
rank * bs, rank * bs + bs - 1, bs, dtype=int).to(self.device)
|
|
|
|
itc_loss = (F.cross_entropy(sim_i2t, targets, label_smoothing=0.1) +
|
|
F.cross_entropy(sim_t2i, targets, label_smoothing=0.1)) / 2
|
|
|
|
# prepare for itm
|
|
text_input_ids_world = torch.cat(dist.all_gather(text_ids))
|
|
text_attention_mask_world = torch.cat(dist.all_gather(text_attn_mask))
|
|
image_embeds_world = torch.cat(dist.all_gather(image_embeds))
|
|
with torch.no_grad():
|
|
weights_t2i = F.softmax(sim_t2i, dim=1) + 1e-4
|
|
weights_t2i[:, rank * bs:rank * bs + bs].fill_diagonal_(0)
|
|
weights_i2t = F.softmax(sim_i2t, dim=1) + 1e-4
|
|
weights_i2t[:, rank * bs:rank * bs + bs].fill_diagonal_(0)
|
|
|
|
# select a negative image for each text
|
|
image_embeds_neg = []
|
|
for b in range(bs):
|
|
neg_idx = torch.multinomial(weights_t2i[b], 1).item()
|
|
image_embeds_neg.append(image_embeds_world[neg_idx])
|
|
image_embeds_neg = torch.stack(image_embeds_neg, dim=0)
|
|
|
|
# select a negative text for each image
|
|
text_ids_neg = []
|
|
text_atts_neg = []
|
|
for b in range(bs):
|
|
neg_idx = torch.multinomial(weights_i2t[b], 1).item()
|
|
text_ids_neg.append(text_input_ids_world[neg_idx])
|
|
text_atts_neg.append(text_attention_mask_world[neg_idx])
|
|
|
|
text_ids_neg = torch.stack(text_ids_neg, dim=0)
|
|
text_atts_neg = torch.stack(text_atts_neg, dim=0)
|
|
|
|
text_ids_all = torch.cat([text_ids, text_ids, text_ids_neg],
|
|
dim=0) # pos, pos, neg
|
|
text_atts_all = torch.cat(
|
|
[text_attn_mask, text_attn_mask, text_atts_neg],
|
|
dim=0,
|
|
)
|
|
|
|
query_tokens_itm = self.query_tokens.expand(text_ids_all.shape[0], -1,
|
|
-1)
|
|
query_atts_itm = torch.ones(
|
|
query_tokens_itm.size()[:-1], dtype=torch.long).to(self.device)
|
|
attention_mask_all = torch.cat([query_atts_itm, text_atts_all], dim=1)
|
|
|
|
image_embeds_all = torch.cat(
|
|
[image_embeds, image_embeds_neg, image_embeds],
|
|
dim=0) # pos, neg, pos
|
|
image_atts_all = torch.ones(
|
|
image_embeds_all.size()[:-1], dtype=torch.long).to(self.device)
|
|
|
|
output_itm = self.multimodal_backbone.bert(
|
|
text_ids_all,
|
|
query_embeds=query_tokens_itm,
|
|
attention_mask=attention_mask_all,
|
|
encoder_hidden_states=image_embeds_all,
|
|
encoder_attention_mask=image_atts_all,
|
|
return_dict=True,
|
|
)
|
|
|
|
vl_embeddings = output_itm.last_hidden_state[:, :query_tokens_itm.
|
|
size(1), :]
|
|
|
|
# create false data samples
|
|
data_samples.extend(
|
|
[DataSample(is_matched=False) for _ in range(2 * bs)])
|
|
loss_multimodal = self.multimodal_head.loss((vl_embeddings, ),
|
|
data_samples)
|
|
|
|
# LM loss
|
|
decoder_input_ids = text_ids.clone()
|
|
decoder_input_ids[:, 0] = self.tokenizer.bos_token_id
|
|
labels = decoder_input_ids.masked_fill(
|
|
decoder_input_ids == self.tokenizer.pad_token_id, -100)
|
|
|
|
query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
|
|
query_atts = torch.ones(
|
|
query_tokens.size()[:-1], dtype=torch.long).to(self.device)
|
|
attention_mask = torch.cat([query_atts, text_attn_mask], dim=1)
|
|
lm_output = self.multimodal_backbone(
|
|
decoder_input_ids,
|
|
attention_mask=attention_mask,
|
|
past_key_values=query_output.past_key_values,
|
|
return_dict=True,
|
|
labels=labels,
|
|
)
|
|
|
|
return dict(
|
|
itc_loss=itc_loss, **loss_multimodal, lm_loss=lm_output.loss)
|
|
|
|
def predict_all(self,
|
|
feats: Dict[str, torch.Tensor],
|
|
data_samples: List[DataSample],
|
|
num_images: int = None,
|
|
num_texts: int = None,
|
|
cal_i2t: bool = True,
|
|
cal_t2i: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Compute similarity matrix between images and texts across all ranks.
|
|
|
|
Args:
|
|
feats (Dict[str, torch.Tensor]): Features from the current rank.
|
|
data_samples (List[DataSample]): Data samples from the current
|
|
rank.
|
|
num_images (int, optional): Number of images to use.
|
|
Defaults to None.
|
|
num_texts (int, optional): Number of texts to use.
|
|
Defaults to None.
|
|
cal_i2t (bool, optional): Whether to compute image-to-text
|
|
similarity. Defaults to True.
|
|
cal_t2i (bool, optional): Whether to compute text-to-image
|
|
similarity. Defaults to True.
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor]: Image-to-text and text-to-image
|
|
similarity matrices.
|
|
"""
|
|
text_ids = feats['text_ids']
|
|
text_attn_mask = feats['text_attn_mask']
|
|
image_embeds = feats.get('image_embeds', None)
|
|
image_feat = feats['image_feat']
|
|
text_feat = feats['text_feat']
|
|
|
|
num_images = num_images or image_feat.size(0)
|
|
num_texts = num_texts or text_feat.size(0)
|
|
|
|
if not self.fast_match:
|
|
image_embeds_all = all_gather_concat(image_embeds)[:num_images]
|
|
else:
|
|
image_embeds_all = None
|
|
image_feat_all = all_gather_concat(image_feat)[:num_images]
|
|
text_feat_all = all_gather_concat(text_feat)[:num_texts]
|
|
text_ids_all = all_gather_concat(text_ids)[:num_texts]
|
|
text_attn_mask_all = all_gather_concat(text_attn_mask)[:num_texts]
|
|
|
|
results = []
|
|
if cal_i2t:
|
|
result_i2t = self.compute_score_matrix_i2t(
|
|
image_feat,
|
|
image_embeds,
|
|
text_feat_all,
|
|
text_ids_all,
|
|
text_attn_mask_all,
|
|
)
|
|
results.append(
|
|
self._get_predictions(result_i2t, data_samples, mode='i2t'))
|
|
if cal_t2i:
|
|
result_t2i = self.compute_score_matrix_t2i(
|
|
image_feat_all,
|
|
image_embeds_all,
|
|
text_feat,
|
|
text_ids,
|
|
text_attn_mask,
|
|
)
|
|
results.append(
|
|
self._get_predictions(result_t2i, data_samples, mode='t2i'))
|
|
return tuple(results)
|
|
|
|
def compute_score_matrix_i2t(self, img_feats: torch.Tensor,
|
|
img_embeds: List[torch.Tensor],
|
|
text_feats: torch.Tensor,
|
|
text_ids: torch.Tensor,
|
|
text_atts: torch.Tensor) -> torch.Tensor:
|
|
"""Compare the score matrix for image-to-text retrieval. Every image
|
|
should compare to all the text features.
|
|
|
|
Args:
|
|
img_feats (torch.Tensor): The input tensor with shape (M, C).
|
|
M stands for numbers of samples on a single GPU.
|
|
img_embeds (List[torch.Tensor]): Image features from each layer of
|
|
the vision backbone.
|
|
text_feats (torch.Tensor): The input tensor with shape (N, C).
|
|
N stands for numbers of all samples on all GPUs.
|
|
text_ids (torch.Tensor): The input tensor with shape (N, C).
|
|
text_atts (torch.Tensor): The input tensor with shape (N, C).
|
|
|
|
Returns:
|
|
torch.Tensor: Score matrix of image-to-text retrieval.
|
|
"""
|
|
|
|
# compute i2t sim matrix
|
|
# TODO: check correctness
|
|
sim_matrix_i2t, _ = (img_feats @ text_feats.t()).max(1)
|
|
if self.fast_match:
|
|
return sim_matrix_i2t
|
|
|
|
score_matrix_i2t = torch.full((img_feats.size(0), text_feats.size(0)),
|
|
-100.0).to(self.device)
|
|
|
|
for i in track_iter_progress(range(img_feats.size(0))):
|
|
sims = sim_matrix_i2t[i]
|
|
topk_sim, topk_idx = sims.topk(k=self.topk, dim=0)
|
|
# get repeated image embeddings
|
|
encoder_output = img_embeds[i].repeat(self.topk, 1, 1)
|
|
encoder_att = torch.ones(
|
|
encoder_output.size()[:-1], dtype=torch.long).to(self.device)
|
|
# query embeds and attention masks
|
|
query_tokens = self.query_tokens.expand(encoder_output.shape[0],
|
|
-1, -1)
|
|
query_atts = torch.ones(
|
|
query_tokens.size()[:-1], dtype=torch.long).to(self.device)
|
|
attention_mask = torch.cat([query_atts, text_atts[topk_idx]],
|
|
dim=1)
|
|
output = self.multimodal_backbone.bert(
|
|
text_ids[topk_idx],
|
|
query_embeds=query_tokens,
|
|
attention_mask=attention_mask,
|
|
encoder_hidden_states=encoder_output,
|
|
encoder_attention_mask=encoder_att,
|
|
return_dict=True,
|
|
)
|
|
score = self.multimodal_head(
|
|
(output.last_hidden_state[:, :query_tokens.size(1), :],
|
|
))[:, :, 1].mean(dim=1)
|
|
score_matrix_i2t[i, topk_idx] = score + topk_sim
|
|
|
|
return score_matrix_i2t
|
|
|
|
def compute_score_matrix_t2i(self, img_feats: torch.Tensor,
|
|
img_embeds: List[torch.Tensor],
|
|
text_feats: torch.Tensor,
|
|
text_ids: torch.Tensor,
|
|
text_atts: torch.Tensor) -> torch.Tensor:
|
|
"""Compare the score matrix for text-to-image retrieval.
|
|
|
|
Every text should compare to all the image features.
|
|
|
|
Args:
|
|
img_feats (torch.Tensor): The input tensor with shape (N, C).
|
|
N stands for numbers of all samples on all GPUs.
|
|
img_embeds (List[torch.Tensor]): Image features from each layer of
|
|
the vision backbone.
|
|
text_feats (torch.Tensor): The input tensor with shape (M, C).
|
|
M stands for numbers of samples on a single GPU.
|
|
text_ids (torch.Tensor): The input tensor with shape (M, C).
|
|
text_atts (torch.Tensor): The input tensor with shape (M, C).
|
|
|
|
Returns:
|
|
torch.Tensor: Score matrix of text-to-image retrieval.
|
|
"""
|
|
|
|
# compute t2i sim matrix
|
|
# TODO: check correctness
|
|
sim_matrix_i2t, _ = (img_feats @ text_feats.t()).max(1)
|
|
sim_matrix_t2i = sim_matrix_i2t.t()
|
|
if self.fast_match:
|
|
return sim_matrix_i2t
|
|
|
|
score_matrix_t2i = torch.full((text_feats.size(0), img_feats.size(0)),
|
|
-100.0).to(self.device)
|
|
|
|
for i in track_iter_progress(range(text_feats.size(0))):
|
|
sims = sim_matrix_t2i[i]
|
|
topk_sim, topk_idx = sims.topk(k=self.topk, dim=0)
|
|
# get topk image embeddings
|
|
encoder_output = img_embeds[topk_idx]
|
|
encoder_att = torch.ones(
|
|
encoder_output.size()[:-1], dtype=torch.long).to(self.device)
|
|
# get query embeds and attention masks
|
|
query_tokens = self.query_tokens.expand(encoder_output.shape[0],
|
|
-1, -1)
|
|
query_atts = torch.ones(
|
|
query_tokens.size()[:-1], dtype=torch.long).to(self.device)
|
|
attention_mask = torch.cat(
|
|
[query_atts, text_atts[i].repeat(self.topk, 1)], dim=1)
|
|
output = self.multimodal_backbone.bert(
|
|
text_ids[i].repeat(self.topk, 1),
|
|
query_embeds=query_tokens,
|
|
attention_mask=attention_mask,
|
|
encoder_hidden_states=encoder_output,
|
|
encoder_attention_mask=encoder_att,
|
|
return_dict=True,
|
|
)
|
|
score = self.multimodal_head(
|
|
(output.last_hidden_state[:, :query_tokens.size(1), :],
|
|
))[:, :, 1].mean(dim=1)
|
|
score_matrix_t2i[i, topk_idx] = score + topk_sim
|
|
|
|
return score_matrix_t2i
|