from collections import OrderedDict import logging import os import torch from torch import nn import torch.nn.functional as F import torch.utils.checkpoint as checkpoint from maskrcnn_benchmark.config import try_to_find from timm.models.layers import DropPath, trunc_normal_ logger = logging.getLogger(__name__) class LayerNorm(nn.Module): def __init__(self, hidden_size, eps=1e-12): """Construct a layernorm module in the TF style (epsilon inside the square root).""" super(LayerNorm, self).__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.bias = nn.Parameter(torch.zeros(hidden_size)) self.variance_epsilon = eps def forward(self, x): pdtype = x.dtype x = x.float() u = x.mean(-1, keepdim=True) s = (x - u).pow(2).mean(-1, keepdim=True) x = (x - u) / torch.sqrt(s + self.variance_epsilon) return self.weight * x.to(pdtype) + self.bias class QuickGELU(nn.Module): def forward(self, x: torch.Tensor): return x * torch.sigmoid(1.702 * x) class ResidualAttentionBlock(nn.Module): def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, drop_path: float = 0.0): super().__init__() self.attn = nn.MultiheadAttention(d_model, n_head) self.ln_1 = LayerNorm(d_model) self.mlp = nn.Sequential( OrderedDict( [ ("c_fc", nn.Linear(d_model, d_model * 4)), ("gelu", QuickGELU()), ("c_proj", nn.Linear(d_model * 4, d_model)), ] ) ) self.ln_2 = LayerNorm(d_model) self.attn_mask = attn_mask self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() def attention(self, x: torch.Tensor, key_padding_mask: torch.Tensor = None): self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask, key_padding_mask=key_padding_mask)[0] def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor = None): x = x + self.drop_path(self.attention(self.ln_1(x), key_padding_mask=key_padding_mask)) x = x + self.drop_path(self.mlp(self.ln_2(x))) return x class CLIPTransformer(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.use_checkpoint = cfg.MODEL.LANGUAGE_BACKBONE.USE_CHECKPOINT print("LANGUAGE BACKBONE USE GRADIENT CHECKPOINTING: ", self.cfg.MODEL.LANGUAGE_BACKBONE.USE_CHECKPOINT) self.context_length = self.cfg.MODEL.CLIP.CONTEXT_LENGTH self.width = self.cfg.MODEL.CLIP.WIDTH self.layers = self.cfg.MODEL.CLIP.LAYERS self.heads = self.cfg.MODEL.CLIP.HEADS self.drop_path = self.cfg.MODEL.CLIP.DROP_PATH self.vocab_size = self.cfg.MODEL.CLIP.VOCAB_SIZE self.token_embedding = nn.Embedding(self.vocab_size, self.width) self.positional_embedding = nn.Parameter(torch.empty(self.context_length, self.width)) # attn_mask = self.build_attention_mask() attn_mask = None dpr = [x.item() for x in torch.linspace(0, self.drop_path, self.layers)] # stochastic depth decay rule self.resblocks = nn.ModuleList( [ResidualAttentionBlock(self.width, self.heads, attn_mask, dpr[i]) for i in range(self.layers)] ) self.ln_final = LayerNorm(self.width) trunc_normal_(self.positional_embedding, std=0.02) # nn.init.normal_(self.token_embedding, std=.02) trunc_normal_(self.token_embedding.weight, std=0.02) self.apply(self._init_weights) # loading pre-trained weight from our CLIP models if len(self.cfg.MODEL.LANGUAGE_BACKBONE.WEIGHT) > 0: self.init_weights(pretrained=try_to_find(self.cfg.MODEL.LANGUAGE_BACKBONE.WEIGHT), pretrained_layers=["*"]) def build_attention_mask(self): # lazily create causal attention mask, with full attention between the vision tokens # pytorch uses additive attention mask; fill with -inf mask = torch.empty(self.context_length, self.context_length) mask.fill_(float("-inf")) mask.triu_(1) # zero out the lower diagonal return mask def _init_weights(self, m): if isinstance(m, (nn.Linear, nn.Conv2d)): trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)): nn.init.constant_(m.bias, 0) def resize_pos_embed_1d(self, posemb, shape_new): # rescale the grid of position embeddings when loading from state_dict ntok_old = posemb.shape[0] if ntok_old > 1: ntok_new = shape_new[0] posemb_grid = posemb.unsqueeze(dim=0).permute(0, 2, 1).unsqueeze(dim=-1) posemb_grid = F.interpolate(posemb_grid, size=[ntok_new, 1], mode="bilinear") posemb_grid = posemb_grid.squeeze(dim=-1).permute(0, 2, 1).squeeze(dim=0) posemb = posemb_grid return posemb def init_weights(self, pretrained="", pretrained_layers=[], verbose=False): if os.path.isfile(pretrained): pretrained_dict = torch.load(pretrained, map_location="cpu") logger.info(f"=> loading pretrained clip text model {pretrained}") model_dict = self.state_dict() need_init_state_dict = {} for k, v in pretrained_dict.items(): need_init = k.split(".")[0] in pretrained_layers or pretrained_layers[0] is "*" if need_init: if k.startswith("text.") and k[5:] in model_dict.keys(): need_init_state_dict[k[5:]] = v # notice the context length now changes from 77 to 256, so we need to resize the positional embedding if "positional_embedding" in need_init_state_dict.keys(): old_pos_embed = need_init_state_dict["positional_embedding"].float() new_pos_embed = self.resize_pos_embed_1d( old_pos_embed, (self.cfg.MODEL.CLIP.CONTEXT_LENGTH, old_pos_embed.shape[1]) ) need_init_state_dict["positional_embedding"] = new_pos_embed self.load_state_dict(need_init_state_dict, strict=True) @torch.jit.ignore def no_weight_decay(self): return { "positional_embedding", "token_embedding", } def forward(self, text): input = text["input_ids"] mask = text["attention_mask"] # get extended attention mask for nn.MultiHeadAttention key_padding_mask = (1.0 - mask).to(torch.bool) x = self.token_embedding(input) # [batch_size, n_ctx, d_model] x = x + self.positional_embedding x = x.permute(1, 0, 2) # NLD -> LND for resblock in self.resblocks: if self.use_checkpoint: x = checkpoint.checkpoint(resblock, x, key_padding_mask) else: x = resblock(x, key_padding_mask) x = x.permute(1, 0, 2) # LND -> NLD x = self.ln_final(x) # x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] ret = {"aggregate": x, "embedded": x, "masks": mask, "hidden": x} return ret