Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from typing import List, Optional | |
| import torch | |
| from mmengine.model import BaseModel | |
| from torch import nn | |
| from mmpretrain.registry import MODELS, TOKENIZER | |
| from mmpretrain.structures import DataSample | |
| class Blip2Caption(BaseModel): | |
| """BLIP2 Caption. | |
| Module for BLIP2 Caption task. | |
| Args: | |
| vision_backbone (dict): The config dict for vision backbone. | |
| text_backbone (dict): The config dict for text backbone. | |
| multimodal_backbone (dict): The config dict for multimodal backbone. | |
| vision_neck (dict): The config dict for vision neck. | |
| 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 | |
| "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. | |
| """ | |
| _no_split_modules = ['BEiTViT', 'OPTDecoderLayer', 'BertLayer'] | |
| def __init__(self, | |
| vision_backbone: dict, | |
| text_backbone: dict, | |
| multimodal_backbone: dict, | |
| vision_neck: 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) -> 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().__init__( | |
| init_cfg=init_cfg, data_preprocessor=data_preprocessor) | |
| self.tokenizer = TOKENIZER.build(tokenizer) | |
| self.eos_token_id = self.tokenizer( | |
| '\n', add_special_tokens=False).input_ids[0] | |
| self.vision_backbone = MODELS.build(vision_backbone) | |
| self.ln_vision_backbone = nn.LayerNorm(self.vision_backbone.embed_dims) | |
| self.vision_neck = MODELS.build(vision_neck) | |
| self.text_backbone = MODELS.build(text_backbone) | |
| self.multimodal_backbone = MODELS.build(multimodal_backbone) | |
| self.multimodal_backbone.cls = None | |
| self.multimodal_backbone.bert.embeddings.word_embeddings = None | |
| self.multimodal_backbone.bert.embeddings.position_embeddings = None | |
| for layer in self.multimodal_backbone.bert.encoder.layer: | |
| layer.output = None | |
| layer.intermediate = None | |
| self.prompt = prompt | |
| self.max_txt_len = max_txt_len | |
| self.num_captions = num_captions | |
| prompt_tokens = self.tokenizer(prompt, return_tensors='pt') | |
| self.prompt_length = prompt_tokens.attention_mask.sum(1) | |
| 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) | |
| # freeze the text backbone | |
| for _, param in self.text_backbone.named_parameters(): | |
| param.requires_grad = False | |
| if hasattr(self, 'register_load_state_dict_post_hook'): | |
| self.register_load_state_dict_post_hook(self._ignore_llm_keys_hook) | |
| def forward( | |
| self, | |
| images: torch.Tensor, | |
| data_samples: Optional[List] = None, | |
| mode: str = 'loss', | |
| ) -> List[DataSample]: | |
| """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): | |
| 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}".') | |
| def predict(self, | |
| images: torch.Tensor, | |
| data_samples: Optional[list] = None, | |
| **kwargs) -> List[DataSample]: | |
| """Predict captions from a batch of inputs. | |
| Args: | |
| images (torch.Tensor): The input 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. | |
| """ | |
| # extract image features from | |
| image_embeds = self.ln_vision_backbone(self.vision_backbone(images)[0]) | |
| image_atts = torch.ones( | |
| image_embeds.size()[:-1], | |
| dtype=torch.long, | |
| ).to(images.device) | |
| # distill image features to query tokens | |
| query_tokens = self.query_tokens.expand(image_embeds.size(0), -1, -1) | |
| query_outputs = self.multimodal_backbone.bert( | |
| query_embeds=query_tokens, | |
| encoder_hidden_states=image_embeds, | |
| encoder_attention_mask=image_atts, | |
| return_dict=True, | |
| ) | |
| inputs_opt = self.vision_neck([query_outputs.last_hidden_state]) | |
| attns_opt = torch.ones( | |
| inputs_opt.size()[:-1], dtype=torch.long).to(images.device) | |
| prompt = [self.prompt] * image_embeds.size(0) | |
| opt_tokens = self.tokenizer( | |
| prompt, return_tensors='pt').to(images.device) | |
| input_ids = opt_tokens.input_ids | |
| attention_mask = torch.cat([attns_opt, opt_tokens.attention_mask], | |
| dim=1) | |
| query_embeds = inputs_opt | |
| outputs = self.text_backbone.generate( | |
| input_ids=input_ids, | |
| query_embeds=query_embeds, | |
| attention_mask=attention_mask, | |
| do_sample=False, | |
| top_p=0.9, | |
| temperature=1., | |
| num_beams=5, | |
| max_new_tokens=self.max_txt_len, | |
| min_length=1, | |
| eos_token_id=self.eos_token_id, | |
| repetition_penalty=1.0, | |
| length_penalty=1.0, | |
| num_return_sequences=self.num_captions, | |
| ) | |
| output_text = self.tokenizer.batch_decode( | |
| outputs[:, self.prompt_length:], skip_special_tokens=True) | |
| output_text = [text.strip() for text in output_text] | |
| out_data_samples = [] | |
| if data_samples is None: | |
| data_samples = [None for _ in range(len(output_text))] | |
| for data_sample, decode_token in zip(data_samples, output_text): | |
| if data_sample is None: | |
| data_sample = DataSample() | |
| data_sample.pred_caption = decode_token | |
| out_data_samples.append(data_sample) | |
| return out_data_samples | |
| def _ignore_llm_keys_hook(module, incompatible_keys): | |
| """Avoid warning missing keys of the LLM model.""" | |
| import re | |
| llm_pattern = '^text_backbone' | |
| for key in list(incompatible_keys.missing_keys): | |
| if re.match(llm_pattern, key): | |
| incompatible_keys.missing_keys.remove(key) | |