# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math from functools import partial from collections import OrderedDict import torch import torch.nn as nn from timm.models.layers import trunc_normal_, lecun_normal_, to_2tuple from timm.models.vision_transformer import Attention from timm.models.layers import Mlp, DropPath from timm.models.helpers import named_apply class Block(nn.Module): def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, ffn_targets=False, return_layer_targets=False): super().__init__() self.norm1 = norm_layer(dim) self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) # specify the targets for feature regression self.ffn_targets = ffn_targets self.return_layer_targets = return_layer_targets def forward(self, x): if isinstance(x, tuple): x = x[0] x = x + self.drop_path(self.attn(self.norm1(x))) ffn_out = self.mlp(self.norm2(x)) x = x + self.drop_path(ffn_out) target = ffn_out if self.ffn_targets else x if self.return_layer_targets: return x, target return x class PatchEmbed(nn.Module): """ 2D Image to Patch Embedding """ def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True): super().__init__() img_size = to_2tuple(img_size) patch_size = to_2tuple(patch_size) self.img_size = img_size self.patch_size = patch_size self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) self.num_patches = self.grid_size[0] * self.grid_size[1] self.flatten = flatten self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() def forward(self, x): B, C, H, W = x.shape patch_H, patch_W = self.patch_size assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}" assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}" x = self.proj(x) if self.flatten: x = x.flatten(2).transpose(1, 2) # BCHW -> BNC x = self.norm(x) return x class VisionTransformer(nn.Module): """ Vision Transformer A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale` - https://arxiv.org/abs/2010.11929 Includes distillation token & head support for `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877 """ def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False, drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None, act_layer=None, weight_init='', ffn_targets=False, return_layer_targets=False): """ Args: img_size (int, tuple): input image size patch_size (int, tuple): patch size in_chans (int): number of input channels num_classes (int): number of classes for classification head embed_dim (int): embedding dimension depth (int): depth of transformer num_heads (int): number of attention heads mlp_ratio (int): ratio of mlp hidden dim to embedding dim qkv_bias (bool): enable bias for qkv if True representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set distilled (bool): model includes a distillation token and head as in DeiT models drop_rate (float): dropout rate attn_drop_rate (float): attention dropout rate drop_path_rate (float): stochastic depth rate embed_layer (nn.Module): patch embedding layer norm_layer: (nn.Module): normalization layer weight_init: (str): weight init scheme ffn_targets (bool): whether we use ffn output or block end as the feature targets return_layer_targets (bool): whether we return every layer targets """ super().__init__() self.num_classes = num_classes self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models self.num_tokens = 2 if distilled else 1 norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) act_layer = act_layer or nn.GELU self.patch_embed = embed_layer( img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) self.pos_drop = nn.Dropout(p=drop_rate) self.ffn_targets = ffn_targets self.return_layer_targets = return_layer_targets dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule self.blocks = nn.Sequential(*[ Block( dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer, ffn_targets=ffn_targets, return_layer_targets=return_layer_targets, ) for i in range(depth)]) self.norm = norm_layer(embed_dim) # Representation layer if representation_size and not distilled: self.num_features = representation_size self.pre_logits = nn.Sequential(OrderedDict([ ('fc', nn.Linear(embed_dim, representation_size)), ('act', nn.Tanh()) ])) else: self.pre_logits = nn.Identity() # Classifier head(s) self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity() self.head_dist = None if distilled: self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity() self.init_weights(weight_init) def init_weights(self, mode=''): assert mode in ('jax', 'jax_nlhb', 'nlhb', '') head_bias = -math.log(self.num_classes) if 'nlhb' in mode else 0. trunc_normal_(self.pos_embed, std=.02) if self.dist_token is not None: trunc_normal_(self.dist_token, std=.02) if mode.startswith('jax'): # leave cls token as zeros to match jax impl named_apply(partial(_init_vit_weights, head_bias=head_bias, jax_impl=True), self) else: trunc_normal_(self.cls_token, std=.02) self.apply(_init_vit_weights) def _init_weights(self, m): # this fn left here for compat with downstream users _init_vit_weights(m) @torch.jit.ignore def no_weight_decay(self): return {'pos_embed', 'cls_token', 'dist_token'} def get_classifier(self): if self.dist_token is None: return self.head else: return self.head, self.head_dist def reset_classifier(self, num_classes, global_pool=''): self.num_classes = num_classes self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() if self.num_tokens == 2: self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity() def forward_features(self, x): x = self.patch_embed(x) cls_token = self.cls_token.expand(x.shape[0], -1, -1) # stole cls_tokens impl from Phil Wang, thanks if self.dist_token is None: x = torch.cat((cls_token, x), dim=1) else: x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1) x = self.pos_drop(x + self.pos_embed) x = self.blocks(x) x = self.norm(x) if self.dist_token is None: return self.pre_logits(x[:, 0]) else: return x[:, 0], x[:, 1] def forward(self, x): x = self.forward_features(x) if self.head_dist is not None: x, x_dist = self.head(x[0]), self.head_dist(x[1]) # x must be a tuple if self.training and not torch.jit.is_scripting(): # during inference, return the average of both classifier predictions return x, x_dist else: return (x + x_dist) / 2 else: x = self.head(x) return x def _init_vit_weights(module: nn.Module, name: str = '', head_bias: float = 0., jax_impl: bool = False): """ ViT weight initialization * When called without n, head_bias, jax_impl args it will behave exactly the same as my original init for compatibility with prev hparam / downstream use cases (ie DeiT). * When called w/ valid n (module name) and jax_impl=True, will (hopefully) match JAX impl """ if isinstance(module, nn.Linear): if name.startswith('head'): nn.init.zeros_(module.weight) nn.init.constant_(module.bias, head_bias) elif name.startswith('pre_logits'): lecun_normal_(module.weight) nn.init.zeros_(module.bias) else: if jax_impl: nn.init.xavier_uniform_(module.weight) if module.bias is not None: if 'mlp' in name: nn.init.normal_(module.bias, std=1e-6) else: nn.init.zeros_(module.bias) else: trunc_normal_(module.weight, std=.02) if module.bias is not None: nn.init.zeros_(module.bias) elif jax_impl and isinstance(module, nn.Conv2d): # NOTE conv was left to pytorch default in my original init lecun_normal_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)): nn.init.zeros_(module.bias) nn.init.ones_(module.weight) def compute_gather_ids(masks): unmask_indices = masks.logical_not().nonzero(as_tuple=False) ids_keep = unmask_indices[:, -1].reshape(masks.shape[0], -1) return ids_keep class MaskedTransformer(VisionTransformer): """Inherit vision transformer from timm""" def __init__(self, mask_style='ibot', **kwargs): super().__init__(**kwargs) assert mask_style in ["ibot", "mae", "none"], "mask_style must be `ibot`, `mae`, or `none`" self.patch_size = self.patch_embed.patch_size if isinstance(self.patch_size, tuple): self.patch_size = self.patch_size[0] nn.init.normal_(self.cls_token, std=1e-6) self.mask_style = mask_style if self.mask_style == "ibot": self.mask_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim)) torch.nn.init.normal_(self.mask_token, std=.02) def interpolate_pos_encoding(self, x, w, h, npatch): previous_dtype = x.dtype N = self.pos_embed.shape[1] - 1 if npatch == N and w == h: return self.pos_embed pos_embed = self.pos_embed.float() class_pos_embed = pos_embed[:, 0] patch_pos_embed = pos_embed[:, 1:] dim = x.shape[-1] w0 = w // self.patch_size h0 = h // self.patch_size # we add a small number to avoid floating point error in the interpolation # see discussion at https://github.com/facebookresearch/dino/issues/8 w0, h0 = w0 + 0.1, h0 + 0.1 patch_pos_embed = nn.functional.interpolate( patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2), scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)), mode="bicubic", ) assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1] patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype) def prepare_tokens_with_masks(self, x, masks=None): """ Args: x: data w/ shape [b, c, h, w] masks: shape [b, n], n is the number of tokens, 1 means masked, 0 means unmasked """ b, c, h, w = x.shape x = self.patch_embed(x) if masks is not None: if self.mask_style == 'ibot': x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype), x) elif self.mask_style == 'mae': # only gather unmasked patches # add pos_embed before shuffle pos_embed = self.interpolate_pos_encoding(x, w, h, npatch=x.shape[1]) x = x + pos_embed[:, 1:, :] ids_keep = compute_gather_ids(masks) x = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, x.shape[-1])) # x = x[masks.logical_not()] # x = x.reshape(b, -1, x.size(-1)) else: raise NotImplementedError(f"mask style {self.mask_style} is not supported") if (masks is None) or (self.mask_style != "mae"): x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1) x = x + self.interpolate_pos_encoding(x, w, h, npatch=x.shape[1]-1) else: # mae-style masking, only need to add cls tokens w/ pos embedding cls_token = self.cls_token + self.pos_embed[:, :1, :] x = torch.cat((cls_token.expand(x.shape[0], -1, -1), x), dim=1) return x def forward_features_list(self, x_list, masks_list): x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)] num_data = len(x) if self.return_layer_targets: all_layer_results = [[] for _ in range(num_data)] for i, blk in enumerate(self.blocks): out = [blk(t) for t in x] x = [o[0] for o in out] # store layer targets for j in range(num_data): all_layer_results[j].append(out[j][1]) all_x = x else: all_x = [self.blocks(t) for t in x] all_layer_results = [None for _ in range(num_data)] output = [] for x, masks, layer_results in zip(all_x, masks_list, all_layer_results): x_norm = self.norm(x) output.append( { "x_norm": x_norm, "x_norm_clstoken": x_norm[:, 0], "x_norm_patchtokens": x_norm[:, 1:], "masks": masks, "layer_results": layer_results, } ) return output def forward_features(self, x, masks=None): if isinstance(x, list): return self.forward_features_list(x, masks) x = self.prepare_tokens_with_masks(x, masks) if self.return_layer_targets: layer_results = [] for i, blk in enumerate(self.blocks): x, lr = blk(x) layer_results.append(lr) else: x = self.blocks(x) layer_results = None x_norm = self.norm(x) return { "x_norm": x_norm, "x_norm_clstoken": x_norm[:, 0], "x_norm_patchtokens": x_norm[:, 1:], "masks": masks, "layer_results": layer_results, } def forward(self, *args, is_training=False, **kwargs): ret = self.forward_features(*args, **kwargs) if is_training: return ret else: return ret["x_norm_clstoken"] def vit_small(patch_size=16, teacher_path=None, **kwargs): model = MaskedTransformer( patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, **kwargs) if teacher_path is not None: checkpoint = torch.load(teacher_path, map_location='cpu') if 'state_dict' in checkpoint: pretrained_dict = checkpoint['state_dict'] elif 'model' in checkpoint: pretrained_dict = checkpoint['model'] else: pretrained_dict = checkpoint pretrained_dict = {k.replace("module.visual.", ""): v for k, v in pretrained_dict.items()} missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False) print('missing_keys: ', missing_keys) print('unexpected_keys: ', unexpected_keys) return model def vit_base(patch_size=16, teacher_path=None, **kwargs): model = MaskedTransformer( patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, **kwargs) if teacher_path is not None: checkpoint = torch.load(teacher_path, map_location='cpu') if 'state_dict' in checkpoint: pretrained_dict = checkpoint['state_dict'] elif 'model' in checkpoint: pretrained_dict = checkpoint['model'] else: pretrained_dict = checkpoint pretrained_dict = {k.replace("module.visual.", ""): v for k, v in pretrained_dict.items()} missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False) print('missing_keys: ', missing_keys) print('unexpected_keys: ', unexpected_keys) return model def vit_large(patch_size=14, teacher_path=None, **kwargs): model = MaskedTransformer( patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16, **kwargs) if teacher_path is not None: checkpoint = torch.load(teacher_path, map_location='cpu') if 'state_dict' in checkpoint: pretrained_dict = checkpoint['state_dict'] elif 'model' in checkpoint: pretrained_dict = checkpoint['model'] else: pretrained_dict = checkpoint pretrained_dict = {k.replace("module.visual.", ""): v for k, v in pretrained_dict.items()} missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False) print('missing_keys: ', missing_keys) print('unexpected_keys: ', unexpected_keys) return model if __name__ == '__main__': import argparse from fvcore.nn import FlopCountAnalysis, parameter_count_table parser = argparse.ArgumentParser(description='PyTorch resnet Training') args = parser.parse_args() with torch.no_grad(): model = vit_base(patch_size=14, num_classes=0, mask_style='ibot') # x = torch.randn(1, 3, 224, 224) # out = model(x) # print(out.shape) print(parameter_count_table(model)) tensor = torch.rand(1, 3, 224, 224) flops = FlopCountAnalysis(model, tensor) print("FLOPs: ", flops.total()/1e9)