# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. # The code in this file is adapted from the BeiT implementation which can be found here: # https://github.com/microsoft/unilm/tree/master/beit import logging import math import numpy as np import random from dataclasses import dataclass, field from typing import Optional from omegaconf import II import torch import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist from fairseq.modules import EMAModule, EMAModuleConfig from fairseq.dataclass import FairseqDataclass from fairseq.models import BaseFairseqModel, register_model logger = logging.getLogger(__name__) @dataclass class Data2VecVisionConfig(FairseqDataclass): layer_scale_init_value: float = field( default=1e-4, metadata={"help": "rescale layer outputs, 0 to disable"} ) num_mask_patches: int = field( default=75, metadata={"help": "number of the visual tokens/patches need be masked"}, ) min_mask_patches_per_block: int = 16 max_mask_patches_per_block: int = 196 image_size: int = 224 patch_size: int = 16 in_channels: int = 3 shared_rel_pos_bias: bool = True drop_path: float = 0.1 attention_dropout: float = 0.0 depth: int = 12 embed_dim: int = 768 num_heads: int = 12 mlp_ratio: int = 4 loss_beta: float = field( default=0, metadata={"help": "beta for smooth l1 loss. 0 means use l2 loss"} ) loss_scale: Optional[float] = field( default=None, metadata={ "help": "scale the reconstruction loss by this constant. if None then scales by 1/sqrt(dim)" }, ) average_top_k_layers: int = field( default=8, metadata={"help": "how many layers to average"} ) end_of_block_targets: bool = True layer_norm_target_layer: bool = False instance_norm_target_layer: bool = False batch_norm_target_layer: bool = False instance_norm_targets: bool = False layer_norm_targets: bool = False ema_decay: float = field(default=0.999, metadata={"help": "initial ema decay rate"}) ema_end_decay: float = field( default=0.9999, metadata={"help": "final ema decay rate"} ) # when to finish annealing ema decay rate ema_anneal_end_step: int = II("optimization.max_update") ema_transformer_only: bool = field( default=True, metadata={"help": "whether to momentum update only the transformer layers"}, ) def get_annealed_rate(start, end, curr_step, total_steps): r = end - start pct_remaining = 1 - curr_step / total_steps return end - r * pct_remaining @register_model("data2vec_vision", dataclass=Data2VecVisionConfig) class Data2VecVisionModel(BaseFairseqModel): def __init__(self, cfg: Data2VecVisionConfig): super().__init__() self.cfg = cfg self.ema = None self.average_top_k_layers = cfg.average_top_k_layers self.loss_beta = cfg.loss_beta self.loss_scale = ( cfg.loss_scale if cfg.loss_scale is not None else 1 / math.sqrt(cfg.embed_dim) ) self.patch_embed = PatchEmbed( img_size=cfg.image_size, patch_size=cfg.patch_size, in_chans=cfg.in_channels, embed_dim=cfg.embed_dim, ) patch_size = self.patch_embed.patch_size self.window_size = ( cfg.image_size // patch_size[0], cfg.image_size // patch_size[1], ) self.cls_emb = nn.Parameter(torch.FloatTensor(1, 1, cfg.embed_dim)) self.mask_emb = nn.Parameter(torch.FloatTensor(1, 1, cfg.embed_dim)) nn.init.trunc_normal_(self.cls_emb, 0.02) nn.init.trunc_normal_(self.mask_emb, 0.02) self.encoder = TransformerEncoder(cfg, self.patch_embed.patch_shape) self.final_proj = nn.Linear(cfg.embed_dim, cfg.embed_dim) self.num_updates = 0 def make_ema_teacher(self): ema_config = EMAModuleConfig( ema_decay=self.cfg.ema_decay, ema_fp32=True, ) self.ema = EMAModule( self.encoder if self.cfg.ema_transformer_only else self, ema_config, ) def set_num_updates(self, num_updates): super().set_num_updates(num_updates) if self.ema is None and self.final_proj is not None: logger.info(f"making ema teacher") self.make_ema_teacher() elif self.training and self.ema is not None: if self.cfg.ema_decay != self.cfg.ema_end_decay: if num_updates >= self.cfg.ema_anneal_end_step: decay = self.cfg.ema_end_decay else: decay = get_annealed_rate( self.cfg.ema_decay, self.cfg.ema_end_decay, num_updates, self.cfg.ema_anneal_end_step, ) self.ema.set_decay(decay) if self.ema.get_decay() < 1: self.ema.step(self.encoder if self.cfg.ema_transformer_only else self) self.num_updates = num_updates def state_dict(self, destination=None, prefix="", keep_vars=False): state = super().state_dict(destination, prefix, keep_vars) if self.ema is not None: state[prefix + "_ema"] = self.ema.fp32_params return state def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): if self.ema is not None: k = prefix + "_ema" assert k in state_dict self.ema.restore(state_dict[k], True) del state_dict[k] return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) @classmethod def build_model(cls, cfg: Data2VecVisionConfig, task=None): """Build a new model instance.""" return cls(cfg) def make_mask(self, bsz, num_masks, min_masks, max_masks): height, width = self.window_size masks = np.zeros(shape=(bsz, height, width), dtype=np.int) for i in range(bsz): mask = masks[i] mask_count = 0 min_aspect = 0.3 max_aspect = 1 / min_aspect log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect)) def _mask(mask, max_mask_patches): delta = 0 for attempt in range(10): target_area = random.uniform(min_masks, max_mask_patches) aspect_ratio = math.exp(random.uniform(*log_aspect_ratio)) h = int(round(math.sqrt(target_area * aspect_ratio))) w = int(round(math.sqrt(target_area / aspect_ratio))) if w < width and h < height: top = random.randint(0, height - h) left = random.randint(0, width - w) num_masked = mask[top : top + h, left : left + w].sum() # Overlap if 0 < h * w - num_masked <= max_mask_patches: for i in range(top, top + h): for j in range(left, left + w): if mask[i, j] == 0: mask[i, j] = 1 delta += 1 if delta > 0: break return delta while mask_count < num_masks: max_mask_patches = min(num_masks - mask_count, max_masks) delta = _mask(mask, max_mask_patches) if delta == 0: break else: mask_count += delta return torch.from_numpy(masks) def forward( self, img, mask: bool = True, layer_results: bool = False, ): x = self.patch_embed(img) batch_size, seq_len, _ = x.size() if mask: mask_indices = self.make_mask( img.size(0), self.cfg.num_mask_patches, self.cfg.min_mask_patches_per_block, self.cfg.max_mask_patches_per_block, ) bool_mask = mask_indices.view(mask_indices.size(0), -1).bool() else: mask_indices = bool_mask = None cls_tokens = self.cls_emb.expand(batch_size, -1, -1) x = torch.cat((cls_tokens, x), dim=1) if self.ema is not None: with torch.no_grad(): self.ema.model.eval() if self.cfg.ema_transformer_only: y = self.ema.model( x, layer_results="end" if self.cfg.end_of_block_targets else "fc", ) else: y = self.ema.model( img, mask=False, layer_results=True, ) y = y[-self.cfg.average_top_k_layers :] permuted = False if self.cfg.instance_norm_target_layer or self.cfg.batch_norm_target_layer: y = [tl.transpose(1, 2) for tl in y] # BTC -> BCT permuted = True if self.cfg.batch_norm_target_layer: y = [ F.batch_norm( tl.float(), running_mean=None, running_var=None, training=True ) for tl in y ] if self.cfg.instance_norm_target_layer: y = [F.instance_norm(tl.float()) for tl in y] if permuted: y = [tl.transpose(1, 2) for tl in y] # BCT -> BTC if self.cfg.layer_norm_target_layer: y = [F.layer_norm(tl.float(), tl.shape[-1:]) for tl in y] y = sum(y) / len(y) if self.cfg.layer_norm_targets: y = F.layer_norm(y.float(), y.shape[-1:]) if self.cfg.instance_norm_targets: y = F.instance_norm(y.float().transpose(1, 2)).transpose(1, 2) y = y[bool_mask].float() if mask_indices is not None: mask_token = self.mask_emb.expand(batch_size, seq_len, -1) w = mask_indices.view(mask_indices.size(0), -1, 1).type_as(mask_token) x[:, 1:] = x[:, 1:] * (1 - w) + mask_token * w if layer_results: enc_layer_results = "end" if self.cfg.end_of_block_targets else "fc" else: enc_layer_results = None x = self.encoder(x, layer_results=enc_layer_results) if layer_results or mask_indices is None: return x x = x[bool_mask].float() if self.loss_beta == 0: loss = F.mse_loss(x, y, reduction="none").sum(dim=-1) else: loss = F.smooth_l1_loss(x, y, reduction="none", beta=self.loss_beta).sum( dim=-1 ) if self.loss_scale > 0: loss = loss * self.loss_scale result = { "losses": {"regression": loss.sum()}, "sample_size": loss.numel(), "target_var": self.compute_var(y), "pred_var": self.compute_var(x), "ema_decay": self.ema.get_decay() * 1000, } return result @staticmethod def compute_var(y): y = y.view(-1, y.size(-1)) if dist.is_initialized(): zc = torch.tensor(y.size(0)).cuda() zs = y.sum(dim=0) zss = (y ** 2).sum(dim=0) dist.all_reduce(zc) dist.all_reduce(zs) dist.all_reduce(zss) var = zss / (zc - 1) - (zs ** 2) / (zc * (zc - 1)) return torch.sqrt(var + 1e-6).mean() else: return torch.sqrt(y.var(dim=0) + 1e-6).mean() def remove_pretraining_modules(self, last_layer=None): self.final_proj = None self.ema = None self.encoder.norm = nn.Identity() self.mask_emb = None if last_layer is not None: self.encoder.layers = nn.ModuleList( l for i, l in enumerate(self.encoder.layers) if i <= last_layer ) class PatchEmbed(nn.Module): """Image to Patch Embedding""" def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): super().__init__() if isinstance(img_size, int): img_size = img_size, img_size if isinstance(patch_size, int): patch_size = patch_size, patch_size num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) self.img_size = img_size self.patch_size = patch_size self.num_patches = num_patches self.conv = nn.Conv2d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size ) def forward(self, x): # BCHW -> BTC x = self.conv(x).flatten(2).transpose(1, 2) return x class Attention(nn.Module): def __init__( self, dim, num_heads=8, qkv_bias=True, attn_drop=0.0, proj_drop=0.0, window_size=None, attn_head_dim=None, ): super().__init__() self.num_heads = num_heads head_dim = dim // num_heads if attn_head_dim is not None: head_dim = attn_head_dim all_head_dim = head_dim * self.num_heads self.scale = head_dim ** -0.5 self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False) if qkv_bias: self.q_bias = nn.Parameter(torch.zeros(all_head_dim)) self.v_bias = nn.Parameter(torch.zeros(all_head_dim)) else: self.q_bias = None self.v_bias = None if window_size: self.window_size = window_size self.num_relative_distance = (2 * window_size[0] - 1) * ( 2 * window_size[1] - 1 ) + 3 self.relative_position_bias_table = nn.Parameter( torch.zeros(self.num_relative_distance, num_heads) ) # 2*Wh-1 * 2*Ww-1, nH # cls to token & token 2 cls & cls to cls # get pair-wise relative position index for each token inside the window coords_h = torch.arange(window_size[0]) coords_w = torch.arange(window_size[1]) coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww relative_coords = ( coords_flatten[:, :, None] - coords_flatten[:, None, :] ) # 2, Wh*Ww, Wh*Ww relative_coords = relative_coords.permute( 1, 2, 0 ).contiguous() # Wh*Ww, Wh*Ww, 2 relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0 relative_coords[:, :, 1] += window_size[1] - 1 relative_coords[:, :, 0] *= 2 * window_size[1] - 1 relative_position_index = torch.zeros( size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype, ) relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww relative_position_index[0, 0:] = self.num_relative_distance - 3 relative_position_index[0:, 0] = self.num_relative_distance - 2 relative_position_index[0, 0] = self.num_relative_distance - 1 self.register_buffer("relative_position_index", relative_position_index) else: self.window_size = None self.relative_position_bias_table = None self.relative_position_index = None self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(all_head_dim, dim) self.proj_drop = nn.Dropout(proj_drop) def forward(self, x, rel_pos_bias=None): B, N, C = x.shape qkv_bias = None if self.q_bias is not None: qkv_bias = torch.cat( ( self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias, ) ) # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias) qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) q, k, v = ( qkv[0], qkv[1], qkv[2], ) # make torchscript happy (cannot use tensor as tuple) q = q * self.scale attn = q @ k.transpose(-2, -1) if self.relative_position_bias_table is not None: assert 1==2 relative_position_bias = self.relative_position_bias_table[ self.relative_position_index.view(-1) ].view( self.window_size[0] * self.window_size[1] + 1, self.window_size[0] * self.window_size[1] + 1, -1, ) # Wh*Ww,Wh*Ww,nH relative_position_bias = relative_position_bias.permute( 2, 0, 1 ).contiguous() # nH, Wh*Ww, Wh*Ww attn = attn + relative_position_bias.unsqueeze(0) print("attn.size() :", attn.size()) print("rel_pos_bias.size() :", rel_pos_bias.size()) if rel_pos_bias is not None: attn = attn + rel_pos_bias attn = attn.softmax(dim=-1) attn = self.attn_drop(attn) x = (attn @ v).transpose(1, 2).reshape(B, N, -1) x = self.proj(x) x = self.proj_drop(x) return x class RelativePositionBias(nn.Module): def __init__(self, window_size, num_heads): super().__init__() self.window_size = window_size self.num_relative_distance = (2 * window_size[0] - 1) * ( 2 * window_size[1] - 1 ) + 3 self.relative_position_bias_table = nn.Parameter( torch.zeros(self.num_relative_distance, num_heads) ) # get pair-wise relative position index for each token inside the window coords_h = torch.arange(window_size[0]) coords_w = torch.arange(window_size[1]) coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww relative_coords = ( coords_flatten[:, :, None] - coords_flatten[:, None, :] ) # 2, Wh*Ww, Wh*Ww relative_coords = relative_coords.permute( 1, 2, 0 ).contiguous() # Wh*Ww, Wh*Ww, 2 relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0 relative_coords[:, :, 1] += window_size[1] - 1 relative_coords[:, :, 0] *= 2 * window_size[1] - 1 relative_position_index = torch.zeros( size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype ) relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww relative_position_index[0, 0:] = self.num_relative_distance - 3 relative_position_index[0:, 0] = self.num_relative_distance - 2 relative_position_index[0, 0] = self.num_relative_distance - 1 self.register_buffer("relative_position_index", relative_position_index) def forward(self): relative_position_bias = self.relative_position_bias_table[ self.relative_position_index.view(-1) ].view( self.window_size[0] * self.window_size[1] + 1, self.window_size[0] * self.window_size[1] + 1, -1, ) # Wh*Ww,Wh*Ww,nH print("self.window_size :", self.window_size) print("self.num_relative_distance :", self.num_relative_distance) print("self.relative_position_index :", self.relative_position_index.size(), self.relative_position_index) print("relative_position_bias.size(), relative_position_bias :",relative_position_bias.size(), relative_position_bias) print("self.relative_position_bias_table.size(), self.relative_position_bias_table :",self.relative_position_bias_table.size(), self.relative_position_bias_table) return relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww class DropPath(nn.Module): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" def __init__(self, drop_prob=None): super(DropPath, self).__init__() self.drop_prob = drop_prob def forward(self, x): if self.drop_prob == 0.0 or not self.training: return x keep_prob = 1 - self.drop_prob shape = (x.shape[0],) + (1,) * ( x.ndim - 1 ) # work with diff dim tensors, not just 2D ConvNets random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) random_tensor.floor_() output = x.div(keep_prob) * random_tensor return output def extra_repr(self) -> str: return "p={}".format(self.drop_prob) class Block(nn.Module): def __init__( self, dim, num_heads, mlp_ratio=4.0, drop=0.0, attn_drop=0.0, drop_path=0.0, init_values=None, window_size=None, ): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = Attention( dim, num_heads=num_heads, attn_drop=attn_drop, proj_drop=drop, window_size=window_size, ) self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = nn.LayerNorm(dim) mlp_hidden_dim = int(dim * mlp_ratio) self.mlp = nn.Sequential( nn.Linear(dim, mlp_hidden_dim), nn.GELU(), nn.Linear(mlp_hidden_dim, dim), nn.Dropout(drop), ) if init_values > 0: self.gamma_1 = nn.Parameter( init_values * torch.ones((dim)), requires_grad=True ) self.gamma_2 = nn.Parameter( init_values * torch.ones((dim)), requires_grad=True ) else: self.gamma_1, self.gamma_2 = None, None def forward(self, x, rel_pos_bias=None): print("inside block :", x.size()) if self.gamma_1 is None: x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias)) fc_feature = self.drop_path(self.mlp(self.norm2(x))) x = x + fc_feature else: x = x + self.drop_path( self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias) ) fc_feature = self.drop_path(self.gamma_2 * self.mlp(self.norm2(x))) x = x + fc_feature return x, fc_feature class TransformerEncoder(nn.Module): def __init__(self, cfg: Data2VecVisionConfig, patch_shape): super().__init__() self.rel_pos_bias = None if cfg.shared_rel_pos_bias: self.rel_pos_bias = RelativePositionBias( window_size=patch_shape, num_heads=cfg.num_heads ) dpr = [ x.item() for x in torch.linspace(0, cfg.drop_path, cfg.depth) ] # stochastic depth decay rule print("TransformerEncoder > patch_shape :", patch_shape) self.blocks = nn.ModuleList( Block( dim=cfg.embed_dim, num_heads=cfg.num_heads, attn_drop=cfg.attention_dropout, drop_path=dpr[i], init_values=cfg.layer_scale_init_value, window_size=patch_shape if not cfg.shared_rel_pos_bias else None, ) for i in range(cfg.depth) ) self.norm = nn.LayerNorm(cfg.embed_dim) self.apply(self.init_weights) self.fix_init_weight() def init_weights(self, m): std = 0.02 if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std=std) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) elif isinstance(m, nn.Conv2d): nn.init.trunc_normal_(m.weight, std=std) if m.bias is not None: nn.init.constant_(m.bias, 0) def fix_init_weight(self): def rescale(param, layer_id): param.div_(math.sqrt(2.0 * layer_id)) for layer_id, layer in enumerate(self.blocks): rescale(layer.attn.proj.weight.data, layer_id + 1) rescale(layer.mlp[2].weight.data, layer_id + 1) def extract_features(self, x, layer_results): rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None z = [] for i, blk in enumerate(self.blocks): x, fc_feature = blk(x, rel_pos_bias=rel_pos_bias) if layer_results == "end": z.append(x) elif layer_results == "fc": z.append(fc_feature) return z if layer_results else self.norm(x) def forward(self, x, layer_results=None): x = self.extract_features(x, layer_results=layer_results) if layer_results: return [z[:, 1:] for z in x] x = x[:, 1:] return x