Source code for mmpretrain.models.multimodal.blip.blip_caption
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional
import torch
from mmengine.model import BaseModel
from mmpretrain.registry import MODELS, TOKENIZER
from mmpretrain.structures import DataSample
[docs]@MODELS.register_module()
class BlipCaption(BaseModel):
"""BLIP Caption.
Args:
vision_encoder (dict): Encoder for extracting image features.
decoder_head (dict): The decoder head module to forward and
calculate loss from processed features.
tokenizer: (Optional[dict]): The config for tokenizer.
Defaults to None.
prompt (str): Prompt used for training and eval.
Defaults to ''.
max_txt_len (int): Max text length of input text.
num_captions (int): Number of captions to be generated for each image.
data_preprocessor (Optional[dict]): The config for preprocessing input
data. If None or no specified type, it will use
"MutimodalDataPreprocessor" as type.
See :class:`MutimodalDataPreprocessor` for more details.
Defaults to None.
init_cfg (Optional[dict]): the config to control the initialization.
Defaults to None.
"""
def __init__(self,
vision_encoder: dict,
decoder_head: dict,
tokenizer: Optional[dict] = None,
prompt: str = '',
max_txt_len: int = 20,
num_captions: int = 1,
data_preprocessor: Optional[dict] = None,
init_cfg: Optional[dict] = None):
if data_preprocessor is None:
data_preprocessor = {}
if isinstance(data_preprocessor, dict):
data_preprocessor.setdefault('type', 'MultiModalDataPreprocessor')
data_preprocessor = MODELS.build(data_preprocessor)
super(BlipCaption, self).__init__(
init_cfg=init_cfg, data_preprocessor=data_preprocessor)
self.tokenizer = TOKENIZER.build(tokenizer)
self.visual_encoder = MODELS.build(vision_encoder)
self.seq_gen_head = MODELS.build(decoder_head)
self.prompt = prompt
self.prompt_length = len(self.tokenizer(self.prompt).input_ids) - 1
self.max_txt_len = max_txt_len
self.num_captions = num_captions
[docs] def forward(
self,
images: torch.Tensor,
data_samples: Optional[List] = None,
mode: str = 'loss',
):
"""The unified entry for a forward process in both training and test.
The method should accept two modes: "predict" and "loss":
- "predict": Forward and return the predictions, which are fully
processed to a list of :obj:`DataSample`.
- "loss": Forward and return a dict of losses according to the given
inputs and data samples.
Note that this method doesn't handle neither back propagation nor
optimizer updating, which are done in the :meth:`train_step`.
Args:
images (torch.Tensor): pre_processed img tensor (N, C, ...).
data_samples (List[DataSample], optional): Data samples with
additional infos.
mode (str): Return what kind of value. Defaults to 'loss'.
Returns:
The return type depends on ``mode``.
- If ``mode="loss"``, return a dict of tensor.
"""
if mode == 'loss':
return self.loss(images, data_samples)
elif mode == 'predict':
return self.predict(images, data_samples)
else:
raise RuntimeError(f'Invalid mode "{mode}".')
[docs] def predict(self, images, data_samples=None, **kwargs):
"""Predict captions from a batch of inputs.
Args:
images (torch.Tensor): The input images tensor with shape
(N, C, ...) in general.
data_samples (List[DataSample], optional): The annotation
data of every samples. Defaults to None.
**kwargs: Other keyword arguments accepted by the ``predict``
method of :attr:`head`.
Returns:
List[DataSample]: Return list of data samples.
"""
# prepare inputs for decoder generation.
image_embeds = self.visual_encoder(images)[0]
image_embeds = torch.repeat_interleave(image_embeds, self.num_captions,
0)
prompt = [self.prompt] * image_embeds.size(0)
prompt = self.tokenizer(
prompt, padding='longest',
return_tensors='pt').to(image_embeds.device)
prompt.input_ids[:, 0] = self.tokenizer.bos_token_id
prompt.input_ids = prompt.input_ids[:, :-1]
decoder_out = self.seq_gen_head.predict(
input_ids=prompt.input_ids,
encoder_hidden_states=image_embeds,
sep_token_id=self.tokenizer.sep_token_id,
pad_token_id=self.tokenizer.pad_token_id,
output_attentions=True,
return_dict_in_generate=True,
)
decode_tokens = self.tokenizer.batch_decode(
decoder_out.sequences, skip_special_tokens=True)
out_data_samples = []
if data_samples is None:
data_samples = [None for _ in range(len(decode_tokens))]
for data_sample, decode_token in zip(data_samples, decode_tokens):
if data_sample is None:
data_sample = DataSample()
data_sample.pred_caption = decode_token[len(self.prompt):]
out_data_samples.append(data_sample)
return out_data_samples
[docs] def loss(self, images, data_samples):
"""Calculate losses from a batch of images and data samples.
Args:
images (torch.Tensor): The input images tensor with shape
(N, C, ...) in general.
data_samples (List[ImageTextDataSample]): The annotation data of
every samples.
Returns:
dict[str, Tensor]: a dictionary of loss components.
"""
image_embeds = self.visual_encoder(images)[0]
raw_text = [self.prompt + ds.gt_caption for ds in data_samples]
text = self.tokenizer(
raw_text,
padding='longest',
truncation=True,
max_length=self.max_txt_len,
return_tensors='pt',
).to(image_embeds.device)
text.input_ids[:, 0] = self.tokenizer.bos_token_id
# prepare targets for forwarding decoder
labels = text.input_ids.masked_fill(
text.input_ids == self.tokenizer.pad_token_id, -100)
labels[:, :self.prompt_length] = -100
# forward decoder
image_atts = torch.ones(
image_embeds.size()[:-1], dtype=torch.long).to(image_embeds.device)
losses = self.seq_gen_head.loss(
input_ids=text.input_ids,
encoder_hidden_states=image_embeds,
encoder_attention_mask=image_atts,
labels=labels,
)
return losses