| import torch |
| from torch import nn |
|
|
| from uniperceiver.config import configurable |
|
|
| from ..layers.transformer_encoder_layer import TransformerEncoderLayer |
| from .build import ENCODER_REGISTRY |
|
|
| import uniperceiver.utils.comm as comm |
|
|
| __all__ = ["StandardViT", "TextEncoder", "VisualEncoder"] |
|
|
|
|
|
|
| @ENCODER_REGISTRY.register() |
| class StandardViT(nn.Module): |
| @configurable |
| def __init__(self, *, num_hidden_layers: int, bert_layers, cfg): |
| super(StandardViT, self).__init__() |
| self.num_hidden_layers = num_hidden_layers |
| self.layers = bert_layers |
| self.cfg = cfg |
| self.name = cfg.NAME |
|
|
| @classmethod |
| def from_config(cls, cfg, global_cfg): |
| if cfg.DROP_PATH_PROB_FIXED: |
| dpr = [cfg.DROP_PATH_PROB for _ in range(cfg.NUM_HIDDEN_LAYERS)] |
| else: |
| dpr = [x.item() for x in torch.linspace(0, cfg.DROP_PATH_PROB, cfg.NUM_HIDDEN_LAYERS)] |
|
|
| layers = [] |
| for i in range(cfg.NUM_HIDDEN_LAYERS): |
| layers.append( |
| TransformerEncoderLayer( |
| d_model=cfg.HIDDEN_SIZE, |
| nhead=cfg.NUM_ATTENTION_HEADS, |
| dim_feedforward=cfg.INTERMEDIATE_SIZE, |
| dropout=cfg.HIDDEN_DROPOUT_PROB, |
| drop_path_ratio=dpr[i], |
| activation=cfg.HIDDEN_ACT, |
| layer_scale=global_cfg.MODEL.LAYER_SCALE, |
| ls_init_values=global_cfg.MODEL.LAYER_SCALE_INIT, |
| batch_first=True, |
| norm_first=True, |
| cfg=cfg, |
| )) |
|
|
| bert_layers = nn.ModuleList( |
| layers |
| ) |
| return { |
| "num_hidden_layers": cfg.NUM_HIDDEN_LAYERS, |
| "bert_layers": bert_layers, |
| "cfg": cfg |
| } |
|
|
| @classmethod |
| def add_config(cls, cfg): |
| pass |
|
|
| def _forward(self, x, attn_mask=None, key_padding_masks=None, history_states=None, *kwargs): |
|
|
| for l, layer_module in enumerate(self.layers): |
| x = layer_module( |
| src=x, src_mask=attn_mask, src_key_padding_mask=key_padding_masks |
| ) |
|
|
| return x |
|
|
|
|
| def forward(self, batched_inputs, return_all=False): |
|
|
| raise NotImplementedError |
|
|
| @ENCODER_REGISTRY.register() |
| class VisualEncoder(StandardViT): |
|
|
| @staticmethod |
| def _construct_attention_masks( data, sample_info, task_info): |
|
|
| return None |
|
|
| def forward(self, data, invalid_mask, sample_info, task_info, **kwargs): |
| |
| |
| attn_mask = self._construct_attention_masks(data, sample_info, task_info) |
| history_states = kwargs.pop('history_states', None) |
| out = self._forward(data, |
| attn_mask, |
| invalid_mask, |
| history_states=history_states, |
| **kwargs, |
| ) |
|
|
| return out |
|
|
|
|
| @ENCODER_REGISTRY.register() |
| class TextEncoder(StandardViT): |
|
|
| @staticmethod |
| def _construct_attention_masks( data, sample_info, task_info): |
| mask_type = torch.bool |
| device = data.device |
|
|
| attn_mask = None |
| if isinstance(sample_info, list): |
| sample_info = sample_info[0] |
| if task_info['task_type'] in ['image_caption', 'video_caption'] and sample_info.get('text_spe_cat', False): |
| total_length = data.shape[1] |
| attn_mask = torch.ones((total_length, total_length), dtype=mask_type, device=device) |
| attn_mask[:total_length // 2, :total_length // 2] = torch.ones( |
| (total_length // 2, total_length // 2), dtype=mask_type, device=device).triu_(diagonal=1) |
| attn_mask[total_length // 2:, : total_length // 2] = torch.ones( |
| (total_length // 2, total_length // 2), |
| dtype=mask_type, |
| device=device).triu_(diagonal=0) |
| attn_mask[total_length // 2:, total_length // 2:] = ~torch.ones( |
| (total_length // 2), |
| dtype=mask_type, |
| device=device).diag() |
|
|
| return attn_mask |
|
|
| def forward(self, data, invalid_mask, sample_info, task_info, **kwargs): |
| |
| |
| attn_mask = self._construct_attention_masks(data, sample_info, task_info) |
| history_states = kwargs.pop('history_states', None) |
| out = self._forward(data, |
| attn_mask, |
| invalid_mask, |
| history_states=history_states, |
| **kwargs) |
|
|
| return out |
|
|
