Shortcuts

Source code for mmpretrain.models.multimodal.flamingo.flamingo

# Copyright (c) OpenMMLab. All rights reserved.
import re
from typing import List, Optional

import torch
from mmengine.model import BaseModel

from mmpretrain.registry import MODELS, TOKENIZER
from mmpretrain.structures import DataSample
from .modules import PerceiverResampler
from .utils import ExtendModule


[docs]@MODELS.register_module() class Flamingo(BaseModel): """The Open Flamingo model for multiple tasks. Args: vision_encoder (dict): The config of the vision encoder. lang_encoder (dict): The config of the language encoder. tokenizer (dict): The tokenizer to encode the text. task (int): The task to perform prediction. zeroshot_prompt (str): Prompt used for zero-shot inference. Defaults to '<image>Output:'. shot_prompt_tmpl (str): Prompt used for few-shot inference. Defaults to ``<image>Output:{caption}<|endofchunk|>``. final_prompt_tmpl (str): Final part of prompt used for inference. Defaults to '<image>Output:'. generation_cfg (dict): The extra generation config, accept the keyword arguments of [~`transformers.GenerationConfig`]. Defaults to an empty dict. 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 (dict, optional): The initialization config. Defaults to None. """ support_tasks = {'caption', 'vqa'} _no_split_modules = [ 'TransformerEncoderLayer', 'PerceiverAttention', 'GatedCrossAttentionBlock', 'FlamingoLayer' ] def __init__( self, vision_encoder: dict, lang_encoder: dict, tokenizer: dict, task: str = 'caption', zeroshot_prompt: str = '<image>Output:', shot_prompt_tmpl: str = '<image>Output:{caption}<|endofchunk|>', final_prompt_tmpl: str = '<image>Output:', generation_cfg: dict = dict(), 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().__init__( init_cfg=init_cfg, data_preprocessor=data_preprocessor) if task not in self.support_tasks: raise ValueError(f'Unsupported task {task}, please select ' f'the task from {self.support_tasks}.') self.task = task # init tokenizer self.tokenizer = TOKENIZER.build(tokenizer) # add Flamingo special tokens to the tokenizer self.tokenizer.add_special_tokens( {'additional_special_tokens': ['<|endofchunk|>', '<image>']}) self.tokenizer.bos_token_id = 1 if self.tokenizer.pad_token is None: # Issue: GPT models don't have a pad token, which we use to # modify labels for the loss. self.tokenizer.add_special_tokens({'pad_token': '<PAD>'}) # Template to format the prompt input self.zeroshot_prompt = zeroshot_prompt self.shot_prompt_tmpl = shot_prompt_tmpl self.final_prompt_tmpl = final_prompt_tmpl # init vision encoder related modules vision_encoder_weight = vision_encoder.pop('pretrained', None) self.vision_encoder = MODELS.build(vision_encoder) if vision_encoder_weight is not None: from mmengine.runner.checkpoint import load_checkpoint load_checkpoint( self.vision_encoder, vision_encoder_weight, map_location='cpu', revise_keys=[(r'^backbone\.', '')], ) self.vision_encoder.is_init = True self.perceiver = PerceiverResampler(dim=self.vision_encoder.embed_dims) # init language encoder related modules self.lang_encoder = ExtendModule(**lang_encoder) self.lang_encoder.resize_token_embeddings(len(self.tokenizer)) self.lang_encoder.media_token_id = self.tokenizer.encode('<image>')[-1] # other necessary parameters self.eoc_token_id = self.tokenizer.encode('<|endofchunk|>')[-1] self.generation_cfg = { 'num_beams': 1, 'max_new_tokens': None, 'temperature': 1.0, 'top_k': 0, 'top_p': 1.0, 'no_repeat_ngram_size': 0, 'prefix_allowed_tokens_fn': None, 'length_penalty': 1.0, 'num_return_sequences': 1, 'do_sample': False, 'early_stopping': False, **generation_cfg, } if hasattr(self, 'register_load_state_dict_post_hook'): self.register_load_state_dict_post_hook(self._load_adapter_hook)
[docs] def forward( self, images: torch.Tensor, data_samples: Optional[List[DataSample]] = None, mode: str = 'loss', ): """The unified entry for a forward process in both training and test. The method should accept only one mode "loss": - "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): The input image tensor with different ndim according to the inputs. data_samples (List[DataSample], optional): The annotation data of every samples. It's required if ``mode="loss"``. Defaults to None. 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 extract_vision_feats(self, images: torch.Tensor) -> torch.Tensor: """Extract vision features. Args: images (torch.Tensor): For zero-shot, the input images tensor is with shape (B, C, H, W), for few-shot, which is (B, T_img, C, H, W) in general. Images in the same chunk are collated along T_img. Video data is not supported yet. Returns: torch.Tensor: Return extracted features. """ if images.ndim == 4: # (B, C, H, W) -> (B, 1, C, H, W) for zero-shot. images = images.unsqueeze(1) b, T = images.shape[:2] # b T c h w -> (b T) c h w images = images.view(b * T, *images.shape[-3:]) with torch.no_grad(): vision_feats = self.vision_encoder(images)[-1][:, 1:] # (b T F) v d -> b T F v d Only support F=1 here vision_feats = vision_feats.view(b, T, 1, *vision_feats.shape[-2:]) vision_feats = self.perceiver(vision_feats) # reshapes to (b, T, n, d) return vision_feats
[docs] def predict(self, images: torch.Tensor, data_samples: Optional[List[DataSample]] = None, **generation_cfg): """Predict generation results from a batch of inputs. Args: images (torch.Tensor): For zero-shot, the input images tensor is with shape (B, C, H, W), for few-shot, which is (B, T_img, C, H, W) in general. Images in the same chunk are collated along T_img. Video data is not supported yet. data_samples (List[DataSample], optional): The annotation data of every samples. Defaults to None. **generation_cfg: Other keyword arguments accepted by the ``generate`` method of :attr:`lang_encoder`. Returns: List[DataSample]: Return list of data samples. """ # generation_cfg in prediction should be dominant generation_cfg = {**self.generation_cfg, **generation_cfg} num_beams = generation_cfg['num_beams'] if num_beams > 1: images = images.repeat_interleave(num_beams, dim=0) # extra vision feats and set as language condition feats vision_x = self.extract_vision_feats(images) for layer in self.lang_encoder._get_decoder_layers(): layer.condition_vis_x(vision_x) input_text = self.preprocess_text(data_samples, device=images.device) outputs = self.lang_encoder.generate( input_text.input_ids, attention_mask=input_text.attention_mask, eos_token_id=self.eoc_token_id, **generation_cfg) # clear conditioned layers for language models self.lang_encoder.clear_conditioned_layers() # remove prefix outputs = outputs[:, len(input_text.input_ids[0]):] return self.post_process(outputs, data_samples)
[docs] def preprocess_text(self, data_samples: List[DataSample], device: torch.device) -> List[DataSample]: """Preprocess text in advance before fed into language model. Args: data_samples (List[DataSample]): The annotation data of every samples. Defaults to None. device (torch.device): Device for text to put on. Returns: List[DataSample]: Return list of data samples. """ prompts = [] for sample in data_samples: if 'shots' in sample: # few-shot shot_prompt = ''.join([ self.shot_prompt_tmpl.format(**shot) for shot in sample.get('shots') ]) else: # zero-shot shot_prompt = self.zeroshot_prompt # add final prompt final_prompt = self.final_prompt_tmpl.format(**sample.to_dict()) prompts.append(shot_prompt + final_prompt) self.tokenizer.padding_side = 'left' input_text = self.tokenizer( prompts, padding='longest', truncation=True, return_tensors='pt', max_length=2000, ).to(device) return input_text
[docs] def post_process( self, outputs: torch.Tensor, data_samples: Optional[List[DataSample]]) -> List[DataSample]: """Perform post process for outputs for different task. Args: outputs (torch.Tensor): The generated outputs. data_samples (List[DataSample], optional): The annotation data of every samples. Returns: List[DataSample]: Return list of data samples. """ outputs = self.tokenizer.batch_decode( outputs, skip_special_tokens=True) if data_samples is None: data_samples = [DataSample() for _ in range(len(outputs))] for output, data_sample in zip(outputs, data_samples): # remove text pattern if self.task == 'caption': data_sample.pred_caption = re.split('Output', output, 1)[0].replace('"', '') elif self.task == 'vqa': data_sample.pred_answer = re.split('Question|Answer', output, 1)[0] return data_samples
@staticmethod def _load_adapter_hook(module, incompatible_keys): """Avoid warning missing keys except adapter keys.""" adapter_patterns = [ '^perceiver', 'lang_encoder.*embed_tokens', 'lang_encoder.*gated_cross_attn_layers', 'lang_encoder.*rotary_emb', ] for key in list(incompatible_keys.missing_keys): if not any(re.match(pattern, key) for pattern in adapter_patterns): incompatible_keys.missing_keys.remove(key) for key in list(incompatible_keys.unexpected_keys): if 'position_ids' in key: incompatible_keys.unexpected_keys.remove(key) if 'lang_encoder.gated_cross_attn_layers' in key: incompatible_keys.unexpected_keys.remove(key)
Read the Docs v: latest
Versions
latest
stable
mmcls-1.x
mmcls-0.x
dev
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.