tennant/MUG_caption
4
1# --------------------------------------------------------2# References:3# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm4# DeiT: https://github.com/facebookresearch/deit5# --------------------------------------------------------6 7from functools import partial8 9import torch10from torch._C import Value11import torch.nn as nn12import numpy as np13 14from timm.models.vision_transformer import PatchEmbed, Block15from transformers import EncoderDecoderModel, BertTokenizer, AutoTokenizer16 17 18from torch import einsum, nn19import torch.nn.functional as F20from einops import rearrange, repeat21 22import torch23import torch.nn as nn24import torch.nn.functional as F25 26class FocalLoss(nn.CrossEntropyLoss):27 ''' Focal loss for classification tasks on imbalanced datasets '''28 29 def __init__(self, gamma=1.0, alpha=None, ignore_index=-100, reduction='none'):30 super().__init__(weight=alpha, ignore_index=ignore_index, reduction='none')31 self.reduction = reduction32 self.gamma = gamma33 34 def forward(self, input_, target):35 cross_entropy = super().forward(input_, target)36 # Temporarily mask out ignore index to '0' for valid gather-indices input.37 # This won't contribute final loss as the cross_entropy contribution38 # for these would be zero.39 target = target * (target != self.ignore_index).long()40 input_prob = torch.gather(F.softmax(input_, 1), 1, target.unsqueeze(1)).squeeze(1)41 loss = torch.pow(1 - input_prob, self.gamma) * cross_entropy42 return torch.mean(loss) if self.reduction == 'mean' \43 else torch.sum(loss) if self.reduction == 'sum' \44 else loss45 46 47# helper functions48 49import math50from functools import reduce51 52def prob_mask_like(t, prob):53 return torch.zeros_like(t).float().uniform_(0, 1) < prob54 55def mask_with_tokens(t, token_ids):56 init_no_mask = torch.full_like(t, False, dtype=torch.bool)57 mask = reduce(lambda acc, el: acc | (t == el), token_ids, init_no_mask)58 return mask59 60def get_mask_subset_with_prob(mask, prob):61 batch, seq_len, device = *mask.shape, mask.device62 max_masked = math.ceil(prob * seq_len)63 64 num_tokens = mask.sum(dim=-1, keepdim=True)65 mask_excess = (mask.cumsum(dim=-1) > (num_tokens * prob).ceil())66 mask_excess = mask_excess[:, :max_masked]67 68 rand = torch.rand((batch, seq_len), device=device).masked_fill(~mask, -1e9)69 _, sampled_indices = rand.topk(max_masked, dim=-1)70 sampled_indices = (sampled_indices + 1).masked_fill_(mask_excess, 0)71 72 new_mask = torch.zeros((batch, seq_len + 1), device=device)73 new_mask.scatter_(-1, sampled_indices, 1)74 return new_mask[:, 1:].bool()75 76 77def exists(val):78 return val is not None79 80def default(val, d):81 return val if exists(val) else d82 83# normalization84# they use layernorm without bias, something that pytorch does not offer85 86 87class LayerNorm(nn.Module):88 def __init__(self, dim):89 super().__init__()90 self.gamma = nn.Parameter(torch.ones(dim))91 self.register_buffer("beta", torch.zeros(dim))92 93 def forward(self, x):94 return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)95 96# residual97class Residual(nn.Module):98 def __init__(self, fn):99 super().__init__()100 self.fn = fn101 102 def forward(self, x, *args, **kwargs):103 return self.fn(x, *args, **kwargs) + x104 105# rotary positional embedding106# https://arxiv.org/abs/2104.09864107class RotaryEmbedding(nn.Module):108 def __init__(self, dim):109 super().__init__()110 inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))111 self.register_buffer("inv_freq", inv_freq)112 113 def forward(self, max_seq_len, *, device):114 seq = torch.arange(max_seq_len, device=device, dtype=self.inv_freq.dtype)115 freqs = einsum("i , j -> i j", seq, self.inv_freq)116 return torch.cat((freqs, freqs), dim=-1)117 118 119def rotate_half(x):120 x = rearrange(x, "... (j d) -> ... j d", j=2)121 x1, x2 = x.unbind(dim=-2)122 return torch.cat((-x2, x1), dim=-1)123 124 125def apply_rotary_pos_emb(pos, t):126 return (t * pos.cos()) + (rotate_half(t) * pos.sin())127 128 129# classic Noam Shazeer paper, except here they use SwiGLU instead of the more popular GELU for gating the feedforward130# https://arxiv.org/abs/2002.05202131class SwiGLU(nn.Module):132 def forward(self, x):133 x, gate = x.chunk(2, dim=-1)134 return F.silu(gate) * x135 136 137# parallel attention and feedforward with residual138# discovered by Wang et al + EleutherAI from GPT-J fame139class ParallelTransformerBlock(nn.Module):140 def __init__(self, dim, dim_head=64, heads=8, ff_mult=4, attn_drop_rate=0.0):141 super().__init__()142 self.norm = LayerNorm(dim)143 144 attn_inner_dim = dim_head * heads145 ff_inner_dim = dim * ff_mult146 self.fused_dims = (attn_inner_dim, dim_head, dim_head, (ff_inner_dim * 2))147 148 self.heads = heads149 self.scale = dim_head**-0.5150 self.rotary_emb = RotaryEmbedding(dim_head)151 152 self.fused_attn_ff_proj = nn.Linear(dim, sum(self.fused_dims), bias=False)153 self.attn_out = nn.Linear(attn_inner_dim, dim, bias=False)154 155 self.ff_out = nn.Sequential(156 SwiGLU(),157 nn.Linear(ff_inner_dim, dim, bias=False)158 )159 160 self.attn_drop_rate = attn_drop_rate161 162 # for caching causal mask and rotary embeddings163 164 self.register_buffer("mask", None, persistent=False)165 self.register_buffer("pos_emb", None, persistent=False)166 167 def get_mask(self, n, device):168 if self.mask is not None and self.mask.shape[-1] >= n:169 return self.mask[:n, :n]170 171 mask = torch.ones((n, n), device=device, dtype=torch.bool).triu(1)172 self.register_buffer("mask", mask, persistent=False)173 return mask174 175 def get_rotary_embedding(self, n, device):176 if self.pos_emb is not None and self.pos_emb.shape[-2] >= n:177 return self.pos_emb[:n]178 179 pos_emb = self.rotary_emb(n, device=device)180 self.register_buffer("pos_emb", pos_emb, persistent=False)181 return pos_emb182 183 def forward(self, x, attn_mask=None):184 """185 Performs self attention and feedforward186 einstein notation187 b - batch188 h - heads189 n, i, j - sequence length (base sequence length, source, target)190 d - feature dimension191 """192 193 n, device, h = x.shape[1], x.device, self.heads194 # pre layernorm195 x = self.norm(x)196 # attention queries, keys, values, and feedforward inner197 q, k, v, ff = self.fused_attn_ff_proj(x).split(self.fused_dims, dim=-1)198 199 # split heads200 # they use multi-query single-key-value attention, yet another Noam Shazeer paper201 # they found no performance loss past a certain scale, and more efficient decoding obviously202 # https://arxiv.org/abs/1911.02150203 q = rearrange(q, "b n (h d) -> b h n d", h=h)204 # rotary embeddings205 positions = self.get_rotary_embedding(n, device)206 q, k = map(lambda t: apply_rotary_pos_emb(positions, t), (q, k))207 # scale208 q = q * self.scale209 # similarity210 sim = einsum("b h i d, b j d -> b h i j", q, k)211 # causal mask212 causal_mask = self.get_mask(n, device)213 sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)214 215 # extra attention mask - for masking out attention from text CLS token to padding216 if exists(attn_mask):217 attn_mask = rearrange(attn_mask, 'b i j -> b 1 i j')218 sim = sim.masked_fill(~attn_mask, -torch.finfo(sim.dtype).max)219 220 if self.attn_drop_rate != 0.:221 # import ipdb; ipdb.set_trace()222 drop_ind = sim != -torch.finfo(sim.dtype).max223 dropout_mask = torch.cuda.FloatTensor(*sim[drop_ind].shape).uniform_() > self.attn_drop_rate224 sim[drop_ind] = sim[drop_ind].masked_fill(~dropout_mask, -torch.finfo(sim.dtype).max)225 226 # attention227 sim = sim - sim.amax(dim=-1, keepdim=True).detach()228 attn = sim.softmax(dim=-1)229 # aggregate values230 out = einsum("b h i j, b j d -> b h i d", attn, v)231 # merge heads232 out = rearrange(out, "b h n d -> b n (h d)")233 return self.attn_out(out) + self.ff_out(ff)234 235# cross attention - using multi-query + one-headed key / values as in PaLM w/ optional parallel feedforward236class CrossAttention(nn.Module):237 def __init__(238 self,239 dim,240 *,241 context_dim=None,242 dim_head=64,243 heads=8,244 parallel_ff=False,245 ff_mult=4,246 norm_context=False,247 dropout=0.0,248 ):249 super().__init__()250 self.heads = heads251 self.scale = dim_head ** -0.5252 inner_dim = heads * dim_head253 context_dim = default(context_dim, dim)254 255 self.norm = LayerNorm(dim)256 self.context_norm = LayerNorm(context_dim) if norm_context else nn.Identity()257 258 self.to_q = nn.Linear(dim, inner_dim, bias=False)259 self.to_kv = nn.Linear(context_dim, dim_head * 2, bias=False)260 self.to_out = nn.Linear(inner_dim, dim, bias=False)261 262 self.dropout = dropout263 264 # whether to have parallel feedforward265 ff_inner_dim = ff_mult * dim266 267 self.ff = nn.Sequential(268 nn.Linear(dim, ff_inner_dim * 2, bias=False),269 SwiGLU(),270 nn.Linear(ff_inner_dim, dim, bias=False)271 ) if parallel_ff else None272 273 def forward(self, x, context):274 """275 Use text and query, and image as kv276 einstein notation277 b - batch278 h - heads279 n, i, j - sequence length (base sequence length, source, target)280 d - feature dimension281 """282 283 # pre-layernorm, for queries and context284 x = self.norm(x)285 context = self.context_norm(context)286 # get queries287 q = self.to_q(x)288 q = rearrange(q, 'b n (h d) -> b h n d', h = self.heads)289 # scale290 q = q * self.scale291 # get key / values292 k, v = self.to_kv(context).chunk(2, dim=-1)293 # query / key similarity294 sim = einsum('b h i d, b j d -> b h i j', q, k)295 296 # dropout297 if self.training:298 dropout_mask = torch.cuda.FloatTensor(*sim.shape).uniform_() > self.dropout299 sim = sim.masked_fill(~dropout_mask, -torch.finfo(sim.dtype).max)300 301 # attention302 sim = sim - sim.amax(dim=-1, keepdim=True)303 attn = sim.softmax(dim=-1)304 # aggregate305 out = einsum('b h i j, b j d -> b h i d', attn, v)306 # merge and combine heads307 out = rearrange(out, 'b h n d -> b n (h d)')308 out = self.to_out(out)309 # add parallel feedforward (for multimodal layers)310 if exists(self.ff):311 out = out + self.ff(x)312 return out313 314 315 316def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):317 """318 grid_size: int of the grid height and width319 return:320 pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)321 """322 grid_h = np.arange(grid_size, dtype=np.float32)323 grid_w = np.arange(grid_size, dtype=np.float32)324 grid = np.meshgrid(grid_w, grid_h) # here w goes first325 grid = np.stack(grid, axis=0)326 327 grid = grid.reshape([2, 1, grid_size, grid_size])328 pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)329 if cls_token:330 pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)331 return pos_embed332 333def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):334 assert embed_dim % 2 == 0335 336 # use half of dimensions to encode grid_h337 emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)338 emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)339 340 emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)341 return emb342 343def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):344 """345 embed_dim: output dimension for each position346 pos: a list of positions to be encoded: size (M,)347 out: (M, D)348 """349 assert embed_dim % 2 == 0350 omega = np.arange(embed_dim // 2, dtype=np.float32)351 omega /= embed_dim / 2.352 omega = 1. / 10000**omega # (D/2,)353 354 pos = pos.reshape(-1) # (M,)355 out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product356 357 emb_sin = np.sin(out) # (M, D/2)358 emb_cos = np.cos(out) # (M, D/2)359 360 emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)361 return emb362 363class MaskedAutoencoderViT(nn.Module):364 """ Masked Autoencoder with VisionTransformer backbone365 """366 def __init__(self, img_size=224, patch_size=16, in_chans=3,367 embed_dim=1024, depth=24, num_heads=16,368 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,369 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=True,370 unimodal_depth=2, multimodal_depth=8, dim_head=64,heads=8,371 ff_mult=4, extract_multi_level=False, use_focal_loss=False, focal_gamma=1.0,372 less_u=False, use_weak_negative=False, use_label_smooth=False, ls_coef=0.1,373 use_maximum_entropy=False, ce_additional=False, use_word_weights=False, use_token_pos=False,374 use_expect_k=False, use_top_k=False, mae_decoder_caption=False, decoder_slot_depth=2, disable_decoder_vis_token_grad=False,375 cross_attn_dropout=0.0, predict_next_k_words=False, next_k=3, masked_text=False, masked_text_ratio=0.25, text_length=70,376 projector_layer=0, uni_dim=1024, uni_dim_head=64, uni_heads=8, uni_ff_mult=4, text_drop_attn=0.):377 super().__init__()378 379 # --------------------------------------------------------------------------380 # MAE encoder specifics381 self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)382 num_patches = self.patch_embed.num_patches383 384 self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))385 self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False) # fixed sin-cos embedding386 387 self.blocks = nn.ModuleList([388 Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)389 for i in range(depth)])390 self.norm = norm_layer(embed_dim)391 # --------------------------------------------------------------------------392 393 # --------------------------------------------------------------------------394 # MAE decoder specifics395 self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)396 397 self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))398 399 self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False) # fixed sin-cos embedding400 401 self.mae_decoder_depth = decoder_depth402 self.mae_decoder_caption = mae_decoder_caption403 self.decoder_blocks = nn.ModuleList([404 Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)405 for i in range(decoder_depth)])406 407 if self.mae_decoder_caption:408 409 self.decoder_slot_layers = nn.ModuleList([])410 for _ in range(decoder_slot_depth):411 self.decoder_slot_layers.append(412 Residual(CrossAttention(dim=decoder_embed_dim, dim_head=dim_head, heads=heads, parallel_ff=True, ff_mult=ff_mult,)),413 # Residual(CrossAttention(dim=decoder_embed_dim, dim_head=dim_head, heads=heads, parallel_ff=True, ff_mult=ff_mult,))414 )415 self.decoder_caption_proj = nn.Linear(decoder_embed_dim, embed_dim)416 self.disable_decoder_vis_token_grad = disable_decoder_vis_token_grad417 418 self.decoder_norm = norm_layer(decoder_embed_dim)419 self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias=True) # encoder to decoder420 # --------------------------------------------------------------------------421 422 self.norm_pix_loss = norm_pix_loss423 424 # --------------------------------------------------------------------------425 # captioner specifics426 # unimodal layer is for text tokens.427 # multimodal layer is for text to query from image. 428 self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased", )429 430 # token embeddings431 # NOTE: +1 for mask token used by MLM objective432 # self.token_emb = nn.Embedding(len(self.tokenizer.vocab) + 1, uni_dim)433 434 self.token_emb = nn.Embedding(len(self.tokenizer.vocab), uni_dim)435 self.text_cls_token = nn.Parameter(torch.randn(uni_dim))436 437 self.embed_dim = embed_dim438 self.uni_dim = uni_dim439 440 #import ipdb; ipdb.set_trace()441 # unimodal layers442 # TODO: search on the four parameters443 # uni_dim=1024, uni_dim_head=64, uni_heads=8, uni_ff_mult=4444 self.text_drop_attn = text_drop_attn445 self.unimodal_layers = nn.ModuleList([])446 for _ in range(unimodal_depth):447 self.unimodal_layers.append(448 Residual(ParallelTransformerBlock(dim=uni_dim, dim_head=uni_dim_head, 449 heads=uni_heads, ff_mult=uni_ff_mult, attn_drop_rate=self.text_drop_attn)),450 )451 452 self.need_uni_2_mul_proj = False453 if uni_dim != embed_dim:454 self.need_uni_2_mul_proj = True455 self.uni_2_mul_proj = nn.Linear(uni_dim, embed_dim)456 457 # multimodal layers458 self.multimodal_layers = nn.ModuleList([])459 self.less_u = less_u460 if less_u:461 for _ in range(multimodal_depth):462 self.multimodal_layers.append(nn.ModuleList([463 Residual(CrossAttention(dim=embed_dim, dim_head=dim_head, heads=heads, parallel_ff=True, ff_mult=ff_mult, dropout=cross_attn_dropout)),464 Residual(CrossAttention(dim=embed_dim, dim_head=dim_head, heads=heads, parallel_ff=True, ff_mult=ff_mult, dropout=cross_attn_dropout))465 ]))466 else:467 for _ in range(multimodal_depth):468 self.multimodal_layers.append(nn.ModuleList([469 Residual(ParallelTransformerBlock(dim=embed_dim, dim_head=dim_head, heads=heads, ff_mult=ff_mult)),470 Residual(CrossAttention(dim=embed_dim, dim_head=dim_head, heads=heads, parallel_ff=True, ff_mult=ff_mult, dropout=cross_attn_dropout))471 ]))472 473 # to logits: for softmax caption loss474 self.to_logits = nn.Sequential(475 LayerNorm(embed_dim),476 nn.Linear(embed_dim, len(self.tokenizer.vocab), bias=False)477 )478 479 self.ce_additional = ce_additional480 if ce_additional:481 # to logits: for other losses482 self.to_logits_1 = nn.Sequential(483 LayerNorm(embed_dim),484 nn.Linear(embed_dim, len(self.tokenizer.vocab), bias=False)485 )486 487 nn.init.normal_(self.token_emb.weight, std=0.02)488 489 self.pad_id = 0490 self.cls_id = 101491 self.sep_id = 102492 493 self.logsoftmax = nn.LogSoftmax(dim=1)494 495 self.extract_multi_level = extract_multi_level496 if self.extract_multi_level:497 self.projectors = nn.ModuleList([nn.Sequential(498 nn.Linear(embed_dim, embed_dim // 2),499 nn.GELU(),500 nn.Linear(embed_dim // 2, embed_dim),501 norm_layer(embed_dim)502 ) for _ in [2, 5, 8,]])503 # --------------------------------------------------------------------------504 505 self.use_focal_loss = use_focal_loss506 507 self.use_weak_negative = use_weak_negative508 self.use_label_smooth = use_label_smooth509 self.ls_coef = ls_coef510 self.use_entropy = use_maximum_entropy511 self.use_word_weights = use_word_weights512 self.use_token_pos = use_token_pos513 514 self.predict_next_k_words = predict_next_k_words515 self.next_k = next_k516 self.pad = torch.nn.ReplicationPad1d((0, self.next_k-1))517 518 self.use_expect_k = use_expect_k519 self.use_top_k = use_top_k520 521 if self.use_word_weights or self.use_token_pos:522 self.focal_loss = FocalLoss(ignore_index=self.pad_id, gamma=focal_gamma, reduction='none')523 else:524 self.focal_loss = FocalLoss(ignore_index=self.pad_id, gamma=focal_gamma, reduction='mean')525 526 self.masked_text = masked_text527 self.masked_text_ratio = masked_text_ratio528 # self.text_mask_token = nn.Parameter(torch.randn(embed_dim))529 self.mask_token_id = len(self.tokenizer.vocab)530 531 # self.text_position_embed = nn.Parameter(torch.zeros(1, text_length, embed_dim), requires_grad=False)532 self.text_length = text_length533 534 self.latent_projector_layer = projector_layer535 if self.latent_projector_layer != 0:536 self.latent_projector = [537 nn.Linear(embed_dim, embed_dim),538 nn.ReLU()539 ] * (self.latent_projector_layer - 1)540 self.latent_projector.append(nn.Linear(embed_dim, embed_dim))541 542 self.latent_projector = nn.Sequential(*self.latent_projector)543 544 545 self.initialize_weights()546 547 548 def initialize_weights(self):549 # initialization550 # initialize (and freeze) pos_embed by sin-cos embedding551 pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)552 self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))553 554 decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)555 self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))556 557 # text_pos_embed = get_1d_sincos_pos_embed_from_grid(self.embed_dim, )558 # torch.nn.init.xavier_normal_(self.text_position_embed) # learnable text position embedding559 560 # initialize patch_embed like nn.Linear (instead of nn.Conv2d)561 w = self.patch_embed.proj.weight.data562 torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))563 564 # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)565 torch.nn.init.normal_(self.cls_token, std=.02)566 torch.nn.init.normal_(self.mask_token, std=.02)567 # torch.nn.init.normal_(self.text_mask_token, std=.02)568 569 # initialize nn.Linear and nn.LayerNorm570 self.apply(self._init_weights)571 572 def _init_weights(self, m):573 if isinstance(m, nn.Linear):574 # we use xavier_uniform following official JAX ViT:575 torch.nn.init.xavier_uniform_(m.weight)576 if isinstance(m, nn.Linear) and m.bias is not None:577 nn.init.constant_(m.bias, 0)578 elif isinstance(m, nn.LayerNorm):579 nn.init.constant_(m.bias, 0)580 nn.init.constant_(m.weight, 1.0)581 582 def patchify(self, imgs):583 """584 imgs: (N, 3, H, W)585 x: (N, L, patch_size**2 *3)586 """587 p = self.patch_embed.patch_size[0]588 assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0589 590 h = w = imgs.shape[2] // p591 x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))592 x = torch.einsum('nchpwq->nhwpqc', x)593 x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))594 return x595 596 def unpatchify(self, x):597 """598 x: (N, L, patch_size**2 *3)599 imgs: (N, 3, H, W)600 """601 p = self.patch_embed.patch_size[0]602 h = w = int(x.shape[1]**.5)603 assert h * w == x.shape[1]604 605 x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))606 x = torch.einsum('nhwpqc->nchpwq', x)607 imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))608 return imgs609 610 def random_masking(self, x, mask_ratio):611 """612 Perform per-sample random masking by per-sample shuffling.613 Per-sample shuffling is done by argsort random noise.614 x: [N, L, D], sequence615 """616 N, L, D = x.shape # batch, length, dim617 len_keep = int(L * (1 - mask_ratio))618 619 noise = torch.rand(N, L, device=x.device) # noise in [0, 1]620 621 # sort noise for each sample622 ids_shuffle = torch.argsort(noise, dim=1) # ascend: small is keep, large is remove623 ids_restore = torch.argsort(ids_shuffle, dim=1)624 625 # keep the first subset626 ids_keep = ids_shuffle[:, :len_keep]627 x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))628 629 # generate the binary mask: 0 is keep, 1 is remove630 mask = torch.ones([N, L], device=x.device)631 mask[:, :len_keep] = 0632 # unshuffle to get the binary mask633 mask = torch.gather(mask, dim=1, index=ids_restore)634 635 return x_masked, mask, ids_restore, ids_keep636 637 def forward_encoder(self, x, mask_ratio):638 # embed patches639 x = self.patch_embed(x)640 641 # add pos embed w/o cls token642 x = x + self.pos_embed[:, 1:, :]643 644 # masking: length -> length * mask_ratio645 x, mask, ids_restore, ids_keep = self.random_masking(x, mask_ratio)646 647 # append cls token648 cls_token = self.cls_token + self.pos_embed[:, :1, :]649 cls_tokens = cls_token.expand(x.shape[0], -1, -1)650 x = torch.cat((cls_tokens, x), dim=1)651 652 if self.extract_multi_level:653 multi_level_feats = []654 # apply Transformer blocks655 for blk_idx, blk in enumerate(self.blocks):656 x = blk(x)657 if blk_idx in [2, 5, 8]:658 multi_level_feats.append(self.projectors[[2,5,8].index(blk_idx)](x))659 x = self.norm(x)660 multi_level_feats.append(x)661 662 return multi_level_feats, mask, ids_restore663 664 665 # apply Transformer blocks666 for blk_idx, blk in enumerate(self.blocks):667 x = blk(x)668 x = self.norm(x)669 670 return x, mask, ids_restore, ids_keep671 672 def forward_decoder(self, x, ids_restore):673 # embed tokens674 x = self.decoder_embed(x)675 # non_mask_token = x676 677 # append mask tokens to sequence678 mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)679 x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1) # no cls token680 x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2])) # unshuffle681 x = torch.cat([x[:, :1, :], x_], dim=1) # append cls token682 683 # add pos embed684 x = x + self.decoder_pos_embed685 686 # apply Transformer blocks687 decoder_feat = []688 for idx, blk in enumerate(self.decoder_blocks):689 x = blk(x)690 if idx == self.mae_decoder_depth // 2:691 decoder_feat.append(x)692 693 x = self.decoder_norm(x)694 695 # use the output from decoder to do captioning696 697 # predictor projection698 x = self.decoder_pred(x)699 700 # remove cls token701 x = x[:, 1:, :]702 703 return x, decoder_feat704 705 def forward_loss(self, imgs, pred, mask):706 """707 imgs: [N, 3, H, W]708 pred: [N, L, p*p*3]709 mask: [N, L], 0 is keep, 1 is remove, 710 """711 target = self.patchify(imgs)712 if self.norm_pix_loss:713 mean = target.mean(dim=-1, keepdim=True)714 var = target.var(dim=-1, keepdim=True)715 target = (target - mean) / (var + 1.e-6)**.5716 717 loss = (pred - target) ** 2718 loss = loss.mean(dim=-1) # [N, L], mean loss per patch719 720 loss = (loss * mask).sum() / mask.sum() # mean loss on removed patches721 return loss722 723 def embed_text(self, text):724 batch, device = text.shape[0], text.device725 726 seq = text.shape[1]727 728 text_tokens = self.token_emb(text)729 730 # append text cls tokens731 text_cls_tokens = repeat(self.text_cls_token, 'd -> b 1 d', b=batch)732 text_tokens = torch.cat((text_tokens, text_cls_tokens), dim=-2)733 734 # create specific mask for text cls token at the end735 # to prevent it from attending to padding736 cls_mask = rearrange(text != self.pad_id, 'b j -> b 1 j')737 attn_mask = F.pad(cls_mask, (0, 1, seq, 0), value=True)738 739 # go through unimodal layers740 for attn_ff in self.unimodal_layers:741 text_tokens = attn_ff(text_tokens, attn_mask=attn_mask)742 743 if self.need_uni_2_mul_proj:744 text_tokens = self.uni_2_mul_proj(text_tokens)745 746 # get text cls token747 text_tokens, text_cls_tokens = text_tokens[:, :-1], text_tokens[:, -1]748 return text_tokens749 750 751 752 def forward(self, imgs, caption_ids=None, attention_mask=None, mask_ratio=0.75, 753 freeze_bert=False, teacher_forcing=False, caption_only=False,754 encoder_only=False, word_weights=None, syn_count=None):755 latent, mask, ids_restore, ids_keep = self.forward_encoder(imgs, mask_ratio)756 757 if not caption_only:758 pred, decoder_feat = self.forward_decoder(latent, ids_restore) # [N, L, p*p*3]759 mae_loss = self.forward_loss(imgs, pred, mask)760 else:761 mae_loss = 0.762 763 if self.latent_projector_layer != 0:764 latent = self.latent_projector(latent)765 766 # latent: visual info: N, L, C767 # caption_ids: N, Len768 text, labels = caption_ids[:, :-1], caption_ids[:, 1:]769 770 seq = text.shape[1]771 text_tokens = self.embed_text(text) # N, Len, C772 773 # create specific mask for text cls token at the end774 # to prevent it from attending to padding775 cls_mask = rearrange(text != self.pad_id, 'b j -> b 1 j')776 attn_mask = F.pad(cls_mask, (0, 1, seq, 0), value=True)777 unimodal_text_tokens = text_tokens778 if not self.less_u:779 for attn_ff, cross_attn in self.multimodal_layers:780 text_tokens = attn_ff(text_tokens, attn_mask=attn_mask[:, :-1, :-1])781 text_tokens = cross_attn(text_tokens, latent)782 else:783 # dim, num_head, 784 for cross_attn1, cross_attn2 in self.multimodal_layers:785 text_tokens = cross_attn1(text_tokens, latent)786 text_tokens = cross_attn2(text_tokens, latent)787 788 logits = self.to_logits(text_tokens) # N, Len, NVocab789 logits = logits.reshape(-1, len(self.tokenizer.vocab))790 labels = labels.reshape(-1)791 792 caption_loss = F.cross_entropy(logits, labels, ignore_index=self.pad_id,)793 794 795 return mae_loss, caption_loss, None796 797 798 799def mae_vit_small_patch16_dec512d8b(**kwargs):800 model = MaskedAutoencoderViT(801 patch_size=16, embed_dim=384, depth=12, num_heads=6,802 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,803 mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)804 return model805 806 807 808def mae_vit_base_patch16_dec512d8b(**kwargs):809 model = MaskedAutoencoderViT(810 patch_size=16, embed_dim=768, depth=12, num_heads=12,811 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,812 mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)813 return model814 815def mae_vit_large_patch16_dec512d8b(**kwargs):816 model = MaskedAutoencoderViT(817 patch_size=16, embed_dim=1024, depth=24, num_heads=16,818 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,819 mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)820 return model821 822 823def mae_vit_huge_patch14_dec512d8b(**kwargs):824 model = MaskedAutoencoderViT(825 patch_size=14, embed_dim=1280, depth=32, num_heads=16,826 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,827 mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)828 return model829 830 831# set recommended archs832mae_vit_small_patch16 = mae_vit_small_patch16_dec512d8b833mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b # decoder: 512 dim, 8 blocks834mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b # decoder: 512 dim, 8 blocks835mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b # decoder: 512 dim, 8 blocks836 837 838 839 840 841