johko/capdec-image-captioning

1from torch import nn2import torch.nn.functional as nnf3from transformers import GPT2Tokenizer, GPT2LMHeadModel4import torch5from typing import Tuple, List, Union, Optional6import numpy as np7 8 9N = type(None)10V = np.array11ARRAY = np.ndarray12ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]13VS = Union[Tuple[V, ...], List[V]]14VN = Union[V, N]15VNS = Union[VS, N]16T = torch.Tensor17TS = Union[Tuple[T, ...], List[T]]18TN = Optional[T]19TNS = Union[Tuple[TN, ...], List[TN]]20TSN = Optional[TS]21TA = Union[T, ARRAY]22 23 24class ClipCaptionModel(nn.Module):25 26    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:27        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)28 29    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor, mask: Optional[torch.Tensor] = None,30                labels: Optional[torch.Tensor] = None):31        embedding_text = self.gpt.transformer.wte(tokens)32        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)33        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)34        if labels is not None:35            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)36            labels = torch.cat((dummy_token, tokens), dim=1)37        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)38        return out39 40    def __init__(self):41        super(ClipCaptionModel, self).__init__()42        self.prefix_length = 4043        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')44        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]45        self.clip_project = TransformerMapper(640, self.gpt_embedding_size, 40,46                                                                     40, 8)47 48 49 50class MLP(nn.Module):51 52    def forward(self, x: T) -> T:53        return self.model(x)54 55    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):56        super(MLP, self).__init__()57        layers = []58        for i in range(len(sizes) -1):59            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))60            if i < len(sizes) - 2:61                layers.append(act())62        self.model = nn.Sequential(*layers)63 64 65class ClipCaptionPrefix(ClipCaptionModel):66 67    def parameters(self, recurse: bool = True):68        return self.clip_project.parameters()69 70    def train(self, mode: bool = True):71        super(ClipCaptionPrefix, self).train(mode)72        self.gpt.eval()73        return self74    75    76class MlpTransformer(nn.Module):77    def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=nnf.relu, dropout=0.):78        super().__init__()79        out_d = out_d if out_d is not None else in_dim80        self.fc1 = nn.Linear(in_dim, h_dim)81        self.act = act82        self.fc2 = nn.Linear(h_dim, out_d)83        self.dropout = nn.Dropout(dropout)84 85    def forward(self, x):86        x = self.fc1(x)87        x = self.act(x)88        x = self.dropout(x)89        x = self.fc2(x)90        x = self.dropout(x)91        return x92 93 94class MultiHeadAttention(nn.Module):95 96    def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.):97        super().__init__()98        self.num_heads = num_heads99        head_dim = dim_self // num_heads100        self.scale = head_dim ** -0.5101        self.to_queries = nn.Linear(dim_self, dim_self, bias=bias)102        self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias)103        self.project = nn.Linear(dim_self, dim_self)104        self.dropout = nn.Dropout(dropout)105 106    def forward(self, x, y=None, mask=None):107        y = y if y is not None else x108        b, n, c = x.shape109        _, m, d = y.shape110        # b n h dh111        queries = self.to_queries(x).reshape(b, n, self.num_heads, c // self.num_heads)112        # b m 2 h dh113        keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads)114        keys, values = keys_values[:, :, 0], keys_values[:, :, 1]115        attention = torch.einsum('bnhd,bmhd->bnmh', queries, keys) * self.scale116        if mask is not None:117            if mask.dim() == 2:118                mask = mask.unsqueeze(1)119            attention = attention.masked_fill(mask.unsqueeze(3), float("-inf"))120        attention = attention.softmax(dim=2)121        out = torch.einsum('bnmh,bmhd->bnhd', attention, values).reshape(b, n, c)122        out = self.project(out)123        return out, attention124 125 126class TransformerLayer(nn.Module):127 128    def forward_with_attention(self, x, y=None, mask=None):129        x_, attention = self.attn(self.norm1(x), y, mask)130        x = x + x_131        x = x + self.mlp(self.norm2(x))132        return x, attention133 134    def forward(self, x, y=None, mask=None):135        x = x + self.attn(self.norm1(x), y, mask)[0]136        x = x + self.mlp(self.norm2(x))137        return x138 139    def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=nnf.relu,140                 norm_layer: nn.Module = nn.LayerNorm):141        super().__init__()142        self.norm1 = norm_layer(dim_self)143        self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout)144        self.norm2 = norm_layer(dim_self)145        self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout)146 147 148class Transformer(nn.Module):149 150    def forward_with_attention(self, x, y=None, mask=None):151        attentions = []152        for layer in self.layers:153            x, att = layer.forward_with_attention(x, y, mask)154            attentions.append(att)155        return x, attentions156 157    def forward(self, x, y=None, mask=None):158        for i, layer in enumerate(self.layers):159            if i % 2 == 0 and self.enc_dec: # cross160                x = layer(x, y)161            elif self.enc_dec:  # self162                x = layer(x, x, mask)163            else:  # self or cross164                x = layer(x, y, mask)165        return x166 167    def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None,168                 mlp_ratio: float = 2., act=nnf.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False):169        super(Transformer, self).__init__()170        dim_ref = dim_ref if dim_ref is not None else dim_self171        self.enc_dec = enc_dec172        if enc_dec:173            num_layers = num_layers * 2174        layers = []175        for i in range(num_layers):176            if i % 2 == 0 and enc_dec:  # cross177                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))178            elif enc_dec:  # self179                layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))180            else:  # self or cross181                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))182        self.layers = nn.ModuleList(layers)183    184 185class TransformerMapper(nn.Module):186 187    def forward(self, x):188        x = self.linear(x).view(x.shape[0], self.clip_length, -1)189        prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)190        prefix = torch.cat((x, prefix), dim=1)191        out = self.transformer(prefix)[:, self.clip_length:]192        return out193 194    def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int, clip_length: int, num_layers: int = 8):195        super(TransformerMapper, self).__init__()196        self.clip_length = clip_length197        self.transformer = Transformer(dim_embedding, 8, num_layers)198        self.linear = nn.Linear(dim_clip, clip_length * dim_embedding)199        self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True)