From 1daa15ecc313fee2fcf281e1ecd71847f00bbd28 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Tue, 4 May 2021 11:19:15 -0700 Subject: [PATCH 1/3] Initial Cait commit. Still some cleanup to do. --- tests/test_models.py | 2 +- timm/models/__init__.py | 1 + timm/models/cait.py | 447 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 449 insertions(+), 1 deletion(-) create mode 100644 timm/models/cait.py diff --git a/tests/test_models.py b/tests/test_models.py index 0d3fde76..9f7106b4 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -15,7 +15,7 @@ if hasattr(torch._C, '_jit_set_profiling_executor'): torch._C._jit_set_profiling_mode(False) # transformer models don't support many of the spatial / feature based model functionalities -NON_STD_FILTERS = ['vit_*', 'tnt_*', 'pit_*', 'swin_*', 'coat_*'] +NON_STD_FILTERS = ['vit_*', 'tnt_*', 'pit_*', 'swin_*', 'coat_*', 'cait_*'] NUM_NON_STD = len(NON_STD_FILTERS) # exclude models that cause specific test failures diff --git a/timm/models/__init__.py b/timm/models/__init__.py index 400e1f64..abf15b52 100644 --- a/timm/models/__init__.py +++ b/timm/models/__init__.py @@ -1,5 +1,6 @@ from .byoanet import * from .byobnet import * +from .cait import * from .coat import * from .cspnet import * from .densenet import * diff --git a/timm/models/cait.py b/timm/models/cait.py new file mode 100644 index 00000000..b82add71 --- /dev/null +++ b/timm/models/cait.py @@ -0,0 +1,447 @@ +# Copyright (c) 2015-present, Facebook, Inc. +# All rights reserved. + +import torch +import torch.nn as nn +from functools import partial + +from .layers import trunc_normal_, DropPath +from .vision_transformer import Mlp, PatchEmbed, _cfg +from .registry import register_model + + +__all__ = ['Cait', 'Class_Attention', 'LayerScale_Block_CA', 'LayerScale_Block', 'Attention_talking_head'] + + +class Class_Attention(nn.Module): + # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py + # with slight modifications to do CA + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.k = nn.Linear(dim, dim, bias=qkv_bias) + self.v = nn.Linear(dim, dim, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x): + B, N, C = x.shape + q = self.q(x[:, 0]).unsqueeze(1).reshape(B, 1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + k = self.k(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + q = q * self.scale + v = self.v(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + attn = (q @ k.transpose(-2, -1)) + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x_cls = (attn @ v).transpose(1, 2).reshape(B, 1, C) + x_cls = self.proj(x_cls) + x_cls = self.proj_drop(x_cls) + + return x_cls + + +class LayerScale_Block_CA(nn.Module): + # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py + # with slight modifications to add CA and LayerScale + def __init__( + self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, attn_block=Class_Attention, + mlp_block=Mlp, init_values=1e-4): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = attn_block( + dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + 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_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + 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) + + def forward(self, x, x_cls): + u = torch.cat((x_cls, x), dim=1) + + x_cls = x_cls + self.drop_path(self.gamma_1 * self.attn(self.norm1(u))) + + x_cls = x_cls + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x_cls))) + + return x_cls + + +class Attention_talking_head(nn.Module): + # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py + # with slight modifications to add Talking Heads Attention (https://arxiv.org/pdf/2003.02436v1.pdf) + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): + super().__init__() + + self.num_heads = num_heads + + head_dim = dim // num_heads + + self.scale = qk_scale or head_dim ** -0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + + self.proj = nn.Linear(dim, dim) + + self.proj_l = nn.Linear(num_heads, num_heads) + self.proj_w = nn.Linear(num_heads, num_heads) + + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0] * self.scale, qkv[1], qkv[2] + + attn = (q @ k.transpose(-2, -1)) + + attn = self.proj_l(attn.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + attn = attn.softmax(dim=-1) + + attn = self.proj_w(attn.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class LayerScale_Block(nn.Module): + # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py + # with slight modifications to add layerScale + def __init__( + self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, attn_block=Attention_talking_head, + mlp_block=Mlp, init_values=1e-4): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = attn_block( + dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + 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_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + 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) + + def forward(self, x): + x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x))) + x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x))) + return x + + +class Cait(nn.Module): + # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py + # with slight modifications to adapt to our cait models + 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=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., + drop_path_rate=0., norm_layer=nn.LayerNorm, global_pool=None, + block_layers=LayerScale_Block, + block_layers_token=LayerScale_Block_CA, + patch_layer=PatchEmbed, act_layer=nn.GELU, + attn_block=Attention_talking_head, mlp_block=Mlp, + init_scale=1e-4, + attn_block_token_only=Class_Attention, + mlp_block_token_only=Mlp, + depth_token_only=2, + mlp_ratio_clstk=4.0): + super().__init__() + + self.num_classes = num_classes + self.num_features = self.embed_dim = embed_dim + + self.patch_embed = patch_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.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) + self.pos_drop = nn.Dropout(p=drop_rate) + + dpr = [drop_path_rate for i in range(depth)] + self.blocks = nn.ModuleList([ + block_layers( + dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, + act_layer=act_layer, attn_block=attn_block, mlp_block=mlp_block, init_values=init_scale) + for i in range(depth)]) + + self.blocks_token_only = nn.ModuleList([ + block_layers_token( + dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio_clstk, qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=0.0, attn_drop=0.0, drop_path=0.0, norm_layer=norm_layer, + act_layer=act_layer, attn_block=attn_block_token_only, + mlp_block=mlp_block_token_only, init_values=init_scale) + for i in range(depth_token_only)]) + + self.norm = norm_layer(embed_dim) + + self.feature_info = [dict(num_chs=embed_dim, reduction=0, module='head')] + self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + trunc_normal_(self.pos_embed, std=.02) + trunc_normal_(self.cls_token, std=.02) + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + 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) + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed', 'cls_token'} + + def forward_features(self, x): + B = x.shape[0] + x = self.patch_embed(x) + + cls_tokens = self.cls_token.expand(B, -1, -1) + + x = x + self.pos_embed + x = self.pos_drop(x) + + for i, blk in enumerate(self.blocks): + x = blk(x) + + for i, blk in enumerate(self.blocks_token_only): + cls_tokens = blk(x, cls_tokens) + + x = torch.cat((cls_tokens, x), dim=1) + + x = self.norm(x) + return x[:, 0] + + def forward(self, x): + x = self.forward_features(x) + x = self.head(x) + + return x + + +@register_model +def cait_xxs24_224(pretrained=False, **kwargs): + model = Cait( + img_size=224, patch_size=16, embed_dim=192, depth=24, num_heads=4, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/XXS24_224.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_xxs24(pretrained=False, **kwargs): + model = Cait( + img_size=384, patch_size=16, embed_dim=192, depth=24, num_heads=4, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/XXS24_384.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_xxs36_224(pretrained=False, **kwargs): + model = Cait( + img_size=224, patch_size=16, embed_dim=192, depth=36, num_heads=4, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/XXS36_224.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_xxs36(pretrained=False, **kwargs): + model = Cait( + img_size=384, patch_size=16, embed_dim=192, depth=36, num_heads=4, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/XXS36_384.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_xs24(pretrained=False, **kwargs): + model = Cait( + img_size=384, patch_size=16, embed_dim=288, depth=24, num_heads=6, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/XS24_384.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_s24_224(pretrained=False, **kwargs): + model = Cait( + img_size=224, patch_size=16, embed_dim=384, depth=24, num_heads=8, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/S24_224.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_s24(pretrained=False, **kwargs): + model = Cait( + img_size=384, patch_size=16, embed_dim=384, depth=24, num_heads=8, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/S24_384.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_s36(pretrained=False, **kwargs): + model = Cait( + img_size=384, patch_size=16, embed_dim=384, depth=36, num_heads=8, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-6, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/S36_384.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_m36(pretrained=False, **kwargs): + model = Cait( + img_size=384, patch_size=16, embed_dim=768, depth=36, num_heads=16, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-6, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/M36_384.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model + + +@register_model +def cait_m48(pretrained=False, **kwargs): + model = Cait( + img_size=448, patch_size=16, embed_dim=768, depth=48, num_heads=16, mlp_ratio=4, qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-6, depth_token_only=2, **kwargs) + + model.default_cfg = _cfg() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url="https://dl.fbaipublicfiles.com/deit/M48_448.pth", + map_location="cpu", check_hash=True + ) + checkpoint_no_module = {} + for k in model.state_dict().keys(): + checkpoint_no_module[k] = checkpoint["model"]['module.' + k] + + model.load_state_dict(checkpoint_no_module) + + return model \ No newline at end of file From 3db12b4b6a7b862698b6fd85ee26b9c924d1c4d3 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Wed, 5 May 2021 17:28:19 -0700 Subject: [PATCH 2/3] Finish CaiT cleanup --- timm/models/cait.py | 331 +++++++++++++++++++------------------------- 1 file changed, 142 insertions(+), 189 deletions(-) diff --git a/timm/models/cait.py b/timm/models/cait.py index b82add71..c16bf86a 100644 --- a/timm/models/cait.py +++ b/timm/models/cait.py @@ -1,19 +1,78 @@ +""" Class-Attention in Image Transformers (CaiT) + +Paper: 'Going deeper with Image Transformers' - https://arxiv.org/abs/2103.17239 + +Original code and weights from https://github.com/facebookresearch/deit, copyright below + +""" # Copyright (c) 2015-present, Facebook, Inc. # All rights reserved. +from copy import deepcopy import torch import torch.nn as nn from functools import partial +from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD +from .helpers import build_model_with_cfg, overlay_external_default_cfg from .layers import trunc_normal_, DropPath -from .vision_transformer import Mlp, PatchEmbed, _cfg +from .vision_transformer import Mlp, PatchEmbed from .registry import register_model -__all__ = ['Cait', 'Class_Attention', 'LayerScale_Block_CA', 'LayerScale_Block', 'Attention_talking_head'] - - -class Class_Attention(nn.Module): +__all__ = ['Cait', 'ClassAttn', 'LayerScaleBlockClassAttn', 'LayerScaleBlock', 'TalkingHeadAttn'] + + +def _cfg(url='', **kwargs): + return { + 'url': url, + 'num_classes': 1000, 'input_size': (3, 384, 384), 'pool_size': None, + 'crop_pct': .9, 'interpolation': 'bicubic', 'fixed_input_size': True, + 'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD, + 'first_conv': 'patch_embed.proj', 'classifier': 'head', + **kwargs + } + + +default_cfgs = dict( + cait_xxs24_224=_cfg( + url='https://dl.fbaipublicfiles.com/deit/XXS24_224.pth', + input_size=(3, 224, 224), + ), + cait_xxs24_384=_cfg( + url='https://dl.fbaipublicfiles.com/deit/XXS24_384.pth', + ), + cait_xxs36_224=_cfg( + url='https://dl.fbaipublicfiles.com/deit/XXS36_224.pth', + input_size=(3, 224, 224), + ), + cait_xxs36_384=_cfg( + url='https://dl.fbaipublicfiles.com/deit/XXS36_384.pth', + ), + cait_xs24_384=_cfg( + url='https://dl.fbaipublicfiles.com/deit/XS24_384.pth', + ), + cait_s24_224=_cfg( + url='https://dl.fbaipublicfiles.com/deit/S24_224.pth', + input_size=(3, 224, 224), + ), + cait_s24_384=_cfg( + url='https://dl.fbaipublicfiles.com/deit/S24_384.pth', + ), + cait_s36_384=_cfg( + url='https://dl.fbaipublicfiles.com/deit/S36_384.pth', + ), + cait_m36_384=_cfg( + url='https://dl.fbaipublicfiles.com/deit/M36_384.pth', + ), + cait_m48_448=_cfg( + url='https://dl.fbaipublicfiles.com/deit/M48_448.pth', + input_size=(3, 448, 448), + ), +) + + +class ClassAttn(nn.Module): # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py # with slight modifications to do CA def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): @@ -48,12 +107,12 @@ class Class_Attention(nn.Module): return x_cls -class LayerScale_Block_CA(nn.Module): +class LayerScaleBlockClassAttn(nn.Module): # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py # with slight modifications to add CA and LayerScale def __init__( self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, attn_block=Class_Attention, + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, attn_block=ClassAttn, mlp_block=Mlp, init_values=1e-4): super().__init__() self.norm1 = norm_layer(dim) @@ -68,15 +127,12 @@ class LayerScale_Block_CA(nn.Module): def forward(self, x, x_cls): u = torch.cat((x_cls, x), dim=1) - x_cls = x_cls + self.drop_path(self.gamma_1 * self.attn(self.norm1(u))) - x_cls = x_cls + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x_cls))) - return x_cls -class Attention_talking_head(nn.Module): +class TalkingHeadAttn(nn.Module): # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py # with slight modifications to add Talking Heads Attention (https://arxiv.org/pdf/2003.02436v1.pdf) def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): @@ -118,12 +174,12 @@ class Attention_talking_head(nn.Module): return x -class LayerScale_Block(nn.Module): +class LayerScaleBlock(nn.Module): # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py # with slight modifications to add layerScale def __init__( self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, attn_block=Attention_talking_head, + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, attn_block=TalkingHeadAttn, mlp_block=Mlp, init_values=1e-4): super().__init__() self.norm1 = norm_layer(dim) @@ -147,17 +203,22 @@ class Cait(nn.Module): # with slight modifications to adapt to our cait models 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=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., - drop_path_rate=0., norm_layer=nn.LayerNorm, global_pool=None, - block_layers=LayerScale_Block, - block_layers_token=LayerScale_Block_CA, - patch_layer=PatchEmbed, act_layer=nn.GELU, - attn_block=Attention_talking_head, mlp_block=Mlp, + num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0., + drop_path_rate=0., + norm_layer=partial(nn.LayerNorm, eps=1e-6), + global_pool=None, + block_layers=LayerScaleBlock, + block_layers_token=LayerScaleBlockClassAttn, + patch_layer=PatchEmbed, + act_layer=nn.GELU, + attn_block=TalkingHeadAttn, + mlp_block=Mlp, init_scale=1e-4, - attn_block_token_only=Class_Attention, + attn_block_token_only=ClassAttn, mlp_block_token_only=Mlp, depth_token_only=2, - mlp_ratio_clstk=4.0): + mlp_ratio_clstk=4.0 + ): super().__init__() self.num_classes = num_classes @@ -237,211 +298,103 @@ class Cait(nn.Module): return x -@register_model -def cait_xxs24_224(pretrained=False, **kwargs): - model = Cait( - img_size=224, patch_size=16, embed_dim=192, depth=24, num_heads=4, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/XXS24_224.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) +def checkpoint_filter_fn(state_dict, model=None): + if 'model' in state_dict: + state_dict = state_dict['model'] + checkpoint_no_module = {} + for k, v in state_dict.items(): + checkpoint_no_module[k.replace('module.', '')] = v + return checkpoint_no_module + + +def _create_cait(variant, pretrained=False, default_cfg=None, **kwargs): + if default_cfg is None: + default_cfg = deepcopy(default_cfgs[variant]) + overlay_external_default_cfg(default_cfg, kwargs) + default_num_classes = default_cfg['num_classes'] + default_img_size = default_cfg['input_size'][-2:] + num_classes = kwargs.pop('num_classes', default_num_classes) + img_size = kwargs.pop('img_size', default_img_size) + + if kwargs.get('features_only', None): + raise RuntimeError('features_only not implemented for Vision Transformer models.') + + model = build_model_with_cfg( + Cait, variant, pretrained, + default_cfg=default_cfg, + img_size=img_size, + num_classes=num_classes, + pretrained_filter_fn=checkpoint_filter_fn, + **kwargs) return model @register_model -def cait_xxs24(pretrained=False, **kwargs): - model = Cait( - img_size=384, patch_size=16, embed_dim=192, depth=24, num_heads=4, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/XXS24_384.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) +def cait_xxs24_224(pretrained=False, **kwargs): + model_args = dict(patch_size=16, embed_dim=192, depth=24, num_heads=4, init_scale=1e-5, **kwargs) + model = _create_cait('cait_xxs24_224', pretrained=pretrained, **model_args) + return model + +@register_model +def cait_xxs24_384(pretrained=False, **kwargs): + model_args = dict(patch_size=16, embed_dim=192, depth=24, num_heads=4, init_scale=1e-5, **kwargs) + model = _create_cait('cait_xxs24_384', pretrained=pretrained, **model_args) return model @register_model def cait_xxs36_224(pretrained=False, **kwargs): - model = Cait( - img_size=224, patch_size=16, embed_dim=192, depth=36, num_heads=4, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/XXS36_224.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) - + model_args = dict(patch_size=16, embed_dim=192, depth=36, num_heads=4, init_scale=1e-5, **kwargs) + model = _create_cait('cait_xxs36_224', pretrained=pretrained, **model_args) return model @register_model -def cait_xxs36(pretrained=False, **kwargs): - model = Cait( - img_size=384, patch_size=16, embed_dim=192, depth=36, num_heads=4, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/XXS36_384.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) - +def cait_xxs36_384(pretrained=False, **kwargs): + model_args = dict(patch_size=16, embed_dim=192, depth=36, num_heads=4, init_scale=1e-5, **kwargs) + model = _create_cait('cait_xxs36_384', pretrained=pretrained, **model_args) return model @register_model -def cait_xs24(pretrained=False, **kwargs): - model = Cait( - img_size=384, patch_size=16, embed_dim=288, depth=24, num_heads=6, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/XS24_384.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) - +def cait_xs24_384(pretrained=False, **kwargs): + model_args = dict(patch_size=16, embed_dim=288, depth=24, num_heads=6, init_scale=1e-5, **kwargs) + model = _create_cait('cait_xs24_384', pretrained=pretrained, **model_args) return model @register_model def cait_s24_224(pretrained=False, **kwargs): - model = Cait( - img_size=224, patch_size=16, embed_dim=384, depth=24, num_heads=8, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/S24_224.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) - + model_args = dict(patch_size=16, embed_dim=384, depth=24, num_heads=8, init_scale=1e-5, **kwargs) + model = _create_cait('cait_s24_224', pretrained=pretrained, **model_args) return model @register_model -def cait_s24(pretrained=False, **kwargs): - model = Cait( - img_size=384, patch_size=16, embed_dim=384, depth=24, num_heads=8, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-5, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/S24_384.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) - +def cait_s24_384(pretrained=False, **kwargs): + model_args = dict(patch_size=16, embed_dim=384, depth=24, num_heads=8, init_scale=1e-5, **kwargs) + model = _create_cait('cait_s24_384', pretrained=pretrained, **model_args) return model @register_model -def cait_s36(pretrained=False, **kwargs): - model = Cait( - img_size=384, patch_size=16, embed_dim=384, depth=36, num_heads=8, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-6, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/S36_384.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) - +def cait_s36_384(pretrained=False, **kwargs): + model_args = dict(patch_size=16, embed_dim=384, depth=36, num_heads=8, init_scale=1e-6, **kwargs) + model = _create_cait('cait_s36_384', pretrained=pretrained, **model_args) return model @register_model -def cait_m36(pretrained=False, **kwargs): - model = Cait( - img_size=384, patch_size=16, embed_dim=768, depth=36, num_heads=16, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-6, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/M36_384.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) - +def cait_m36_384(pretrained=False, **kwargs): + model_args = dict(patch_size=16, embed_dim=768, depth=36, num_heads=16, init_scale=1e-6, **kwargs) + model = _create_cait('cait_m36_384', pretrained=pretrained, **model_args) return model @register_model -def cait_m48(pretrained=False, **kwargs): - model = Cait( - img_size=448, patch_size=16, embed_dim=768, depth=48, num_heads=16, mlp_ratio=4, qkv_bias=True, - norm_layer=partial(nn.LayerNorm, eps=1e-6), init_scale=1e-6, depth_token_only=2, **kwargs) - - model.default_cfg = _cfg() - if pretrained: - checkpoint = torch.hub.load_state_dict_from_url( - url="https://dl.fbaipublicfiles.com/deit/M48_448.pth", - map_location="cpu", check_hash=True - ) - checkpoint_no_module = {} - for k in model.state_dict().keys(): - checkpoint_no_module[k] = checkpoint["model"]['module.' + k] - - model.load_state_dict(checkpoint_no_module) - - return model \ No newline at end of file +def cait_m48_448(pretrained=False, **kwargs): + model_args = dict(patch_size=16, embed_dim=768, depth=48, num_heads=16, init_scale=1e-6, **kwargs) + model = _create_cait('cait_m48_448', pretrained=pretrained, **model_args) + return model From d45e50b9db3745c1561dfae1813401576f60e336 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Wed, 5 May 2021 17:51:23 -0700 Subject: [PATCH 3/3] Update test for cait 448x448 model --- tests/test_models.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index 96c51589..ced2fd76 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -43,7 +43,9 @@ def test_model_forward(model_name, batch_size): input_size = model.default_cfg['input_size'] if any([x > MAX_FWD_SIZE for x in input_size]): - # cap forward test at max res 448 * 448 to keep resource down + if is_model_default_key(model_name, 'fixed_input_size'): + pytest.skip("Fixed input size model > limit.") + # cap forward test at max res 384 * 384 to keep resource down input_size = tuple([min(x, MAX_FWD_SIZE) for x in input_size]) inputs = torch.randn((batch_size, *input_size)) outputs = model(inputs)