import math import torch import torch.nn as nn from timm.models.helpers import load_pretrained from timm.models.registry import register_model from timm.models.layers import trunc_normal_ import numpy as np from .MBConv import MBConvBlock from .MHSA import MHSABlock,Mlp def _cfg(url='', **kwargs): return { 'url': url, 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None, 'crop_pct': .9, 'interpolation': 'bicubic', 'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225), 'classifier': 'head', **kwargs } default_cfgs = { 'MetaFG_0': _cfg(), 'MetaFG_1': _cfg(), 'MetaFG_2': _cfg(), } def make_blocks(stage_index,depths,embed_dims,img_size,dpr,extra_token_num=1,num_heads=8,mlp_ratio=4.,stage_type='conv'): stage_name = f'stage_{stage_index}' blocks = [] for block_idx in range(depths[stage_index]): stride = 2 if block_idx == 0 and stage_index != 1 else 1 in_chans = embed_dims[stage_index] if block_idx != 0 else embed_dims[stage_index-1] out_chans = embed_dims[stage_index] image_size = img_size if block_idx == 0 or stage_index == 1 else img_size//2 drop_path_rate = dpr[sum(depths[1:stage_index])+block_idx] if stage_type == 'conv': blocks.append(MBConvBlock(ksize=3,input_filters=in_chans,output_filters=out_chans, image_size=image_size,expand_ratio=int(mlp_ratio),stride=stride,drop_connect_rate=drop_path_rate)) elif stage_type == 'mhsa': blocks.append(MHSABlock(input_dim=in_chans,output_dim=out_chans, image_size=image_size,stride=stride,num_heads=num_heads,extra_token_num=extra_token_num, mlp_ratio=mlp_ratio,drop_path=drop_path_rate)) else: raise NotImplementedError("We only support conv and mhsa") return blocks class MetaFG(nn.Module): def __init__(self,img_size=224,in_chans=3, num_classes=1000, conv_embed_dims = [64,96,192],attn_embed_dims=[384,768], conv_depths = [2,2,3],attn_depths = [5,2],num_heads=32,extra_token_num=1,mlp_ratio=4., conv_norm_layer=nn.BatchNorm2d,attn_norm_layer=nn.LayerNorm, conv_act_layer=nn.ReLU,attn_act_layer=nn.GELU, qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,drop_path_rate=0., meta_dims=[], only_last_cls=False, use_checkpoint=False): super().__init__() self.only_last_cls = only_last_cls self.img_size = img_size self.num_classes = num_classes stem_chs = (3 * (conv_embed_dims[0] // 4), conv_embed_dims[0]) dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(conv_depths[1:]+attn_depths))] #stage_0 self.stage_0 = nn.Sequential(*[ nn.Conv2d(in_chans, stem_chs[0], 3, stride=2, padding=1, bias=False), conv_norm_layer(stem_chs[0]), conv_act_layer(inplace=True), nn.Conv2d(stem_chs[0], stem_chs[1], 3, stride=1, padding=1, bias=False), conv_norm_layer(stem_chs[1]), conv_act_layer(inplace=True), nn.Conv2d(stem_chs[1], conv_embed_dims[0], 3, stride=1, padding=1, bias=False)]) self.bn1 = conv_norm_layer(conv_embed_dims[0]) self.act1 = conv_act_layer(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) #stage_1 self.stage_1 = nn.ModuleList(make_blocks(1,conv_depths+attn_depths,conv_embed_dims+attn_embed_dims,img_size//4, dpr=dpr,num_heads=num_heads,extra_token_num=extra_token_num,mlp_ratio=mlp_ratio,stage_type='conv')) #stage_2 self.stage_2 = nn.ModuleList(make_blocks(2,conv_depths+attn_depths,conv_embed_dims+attn_embed_dims,img_size//4, dpr=dpr,num_heads=num_heads,extra_token_num=extra_token_num,mlp_ratio=mlp_ratio,stage_type='conv')) #stage_3 self.cls_token_1 = nn.Parameter(torch.zeros(1, 1, attn_embed_dims[0])) self.stage_3 = nn.ModuleList(make_blocks(3,conv_depths+attn_depths,conv_embed_dims+attn_embed_dims,img_size//8, dpr=dpr,num_heads=num_heads,extra_token_num=extra_token_num,mlp_ratio=mlp_ratio,stage_type='mhsa')) #stage_4 self.cls_token_2 = nn.Parameter(torch.zeros(1, 1, attn_embed_dims[1])) self.stage_4 = nn.ModuleList(make_blocks(4,conv_depths+attn_depths,conv_embed_dims+attn_embed_dims,img_size//16, dpr=dpr,num_heads=num_heads,extra_token_num=extra_token_num,mlp_ratio=mlp_ratio,stage_type='mhsa')) self.norm_2 = attn_norm_layer(attn_embed_dims[1]) #Aggregate if not self.only_last_cls: self.cl_1_fc = nn.Sequential(*[Mlp(in_features=attn_embed_dims[0], out_features=attn_embed_dims[1]), attn_norm_layer(attn_embed_dims[1])]) self.aggregate = torch.nn.Conv1d(in_channels=2, out_channels=1, kernel_size=1) self.norm_1 = attn_norm_layer(attn_embed_dims[0]) self.norm = attn_norm_layer(attn_embed_dims[1]) # Classifier head self.head = nn.Linear(attn_embed_dims[-1], num_classes) if num_classes > 0 else nn.Identity() trunc_normal_(self.cls_token_1, std=.02) trunc_normal_(self.cls_token_2, 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) elif isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') # fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels # fan_out //= m.groups # m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) # if m.bias is not None: # m.bias.data.zero_() elif isinstance(m, nn.BatchNorm2d): nn.init.ones_(m.weight) nn.init.zeros_(m.bias) @torch.jit.ignore def no_weight_decay(self): return {'cls_token_1','cls_token_2'} def get_classifier(self): return self.head 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() def forward_features(self,x,meta=None): extra_tokens_1 = [self.cls_token_1] extra_tokens_2 = [self.cls_token_2] B = x.shape[0] x = self.stage_0(x) x = self.bn1(x) x = self.act1(x) x = self.maxpool(x) for blk in self.stage_1: x = blk(x) for blk in self.stage_2: x = blk(x) H0,W0 = self.img_size//8,self.img_size//8 for ind,blk in enumerate(self.stage_3): if ind==0: x = blk(x,H0,W0,extra_tokens_1) else: x = blk(x,H0,W0) if not self.only_last_cls: cls_1 = x[:, :1, :] cls_1 = self.norm_1(cls_1) cls_1 = self.cl_1_fc(cls_1) x = x[:, 1:, :] H1,W1 = self.img_size//16,self.img_size//16 x = x.reshape(B,H1,W1,-1).permute(0, 3, 1, 2).contiguous() for ind,blk in enumerate(self.stage_4): if ind==0: x = blk(x,H1,W1,extra_tokens_2) else: x = blk(x,H1,W1) cls_2 = x[:, :1, :] cls_2 = self.norm_2(cls_2) if not self.only_last_cls: cls = torch.cat((cls_1,cls_2), dim=1)#B,2,C cls = self.aggregate(cls).squeeze(dim=1)#B,C cls = self.norm(cls) else: cls = cls_2.squeeze(dim=1) return cls def forward(self, x,meta=None): x = self.forward_features(x,meta) x = self.head(x) return x @register_model def MetaFG_0(pretrained=False, **kwargs): model = MetaFG(conv_embed_dims = [64,96,192],attn_embed_dims=[384,768], conv_depths = [2,2,3],attn_depths = [5,2],num_heads=8,mlp_ratio=4., **kwargs) model.default_cfg = default_cfgs['MetaFG_0'] if pretrained: load_pretrained( model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) return model @register_model def MetaFG_1(pretrained=False, **kwargs): model = MetaFG(conv_embed_dims = [64,96,192],attn_embed_dims=[384,768], conv_depths = [2,2,6],attn_depths = [14,2],num_heads=8,mlp_ratio=4., **kwargs) model.default_cfg = default_cfgs['MetaFG_1'] if pretrained: load_pretrained( model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) return model @register_model def MetaFG_2(pretrained=False, **kwargs): model = MetaFG(conv_embed_dims = [128,128,256],attn_embed_dims=[512,1024], conv_depths = [2,2,6],attn_depths = [14,2],num_heads=8,mlp_ratio=4., **kwargs) model.default_cfg = default_cfgs['MetaFG_2'] if pretrained: load_pretrained( model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3)) return model if __name__ == "__main__": x = torch.randn([2, 3, 224, 224]) model = MetaFG() import ipdb;ipdb.set_trace() output = model(x) print(output.shape)