# https://github.com/mmaaz60/EdgeNeXt # https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/edgenext.py from src.models.bb.edgenext.edgenext_blocks import * import torch from torch import nn from timm.models.layers import trunc_normal_ class EdgeNeXtBNHS(nn.Module): def __init__(self, in_chans=3, num_classes=1000, depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], global_block=[0, 0, 0, 3], global_block_type=['None', 'None', 'None', 'SDTA_BN_HS'], drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., expan_ratio=4, kernel_sizes=[7, 7, 7, 7], heads=[8, 8, 8, 8], use_pos_embd_xca=[False, False, False, False], use_pos_embd_global=False, d2_scales=[2, 3, 4, 5], **kwargs): super().__init__() for g in global_block_type: assert g in ['None', 'SDTA_BN_HS'] if use_pos_embd_global: self.pos_embd = PositionalEncodingFourier(dim=dims[0]) else: self.pos_embd = None self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers stem = nn.Sequential( nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4, bias=False), nn.BatchNorm2d(dims[0]) ) self.downsample_layers.append(stem) for i in range(3): downsample_layer = nn.Sequential( nn.BatchNorm2d(dims[i]), nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2, bias=False), ) self.downsample_layers.append(downsample_layer) self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] cur = 0 for i in range(4): stage_blocks = [] for j in range(depths[i]): if j > depths[i] - global_block[i] - 1: if global_block_type[i] == 'SDTA_BN_HS': stage_blocks.append(SDTAEncoderBNHS(dim=dims[i], drop_path=dp_rates[cur + j], expan_ratio=expan_ratio, scales=d2_scales[i], use_pos_emb=use_pos_embd_xca[i], num_heads=heads[i], )) else: raise NotImplementedError else: stage_blocks.append(ConvEncoderBNHS(dim=dims[i], drop_path=dp_rates[cur + j], layer_scale_init_value=layer_scale_init_value, expan_ratio=expan_ratio, kernel_size=kernel_sizes[i])) self.stages.append(nn.Sequential(*stage_blocks)) cur += depths[i] self.norm = nn.BatchNorm2d(dims[-1]) self.head = nn.Linear(dims[-1], num_classes) self.apply(self._init_weights) self.head_dropout = nn.Dropout(kwargs["classifier_dropout"]) self.head.weight.data.mul_(head_init_scale) self.head.bias.data.mul_(head_init_scale) def _init_weights(self, m): # TODO: MobileViT is using 'kaiming_normal' for initializing conv layers if isinstance(m, (nn.Conv2d, nn.Linear)): trunc_normal_(m.weight, std=.02) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, (LayerNorm, nn.LayerNorm)): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def forward_features(self, x): x = self.downsample_layers[0](x) x = self.stages[0](x) if self.pos_embd: B, C, H, W = x.shape x = x + self.pos_embd(B, H, W) for i in range(1, 4): x = self.downsample_layers[i](x) x = self.stages[i](x) return self.norm(x).mean([-2, -1]) def forward(self, x): x = self.forward_features(x) x = self.head(self.head_dropout(x)) return class EdgeNeXt(nn.Module): def __init__(self, in_chans=3, num_classes=1000, depths=[3, 3, 9, 3], dims=[24, 48, 88, 168], global_block=[0, 0, 0, 3], global_block_type=['None', 'None', 'None', 'SDTA'], drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., expan_ratio=4, kernel_sizes=[7, 7, 7, 7], heads=[8, 8, 8, 8], use_pos_embd_xca=[False, False, False, False], use_pos_embd_global=False, d2_scales=[2, 3, 4, 5], class_drop_rate=0.2, **kwargs): super().__init__() for g in global_block_type: assert g in ['None', 'SDTA'] if use_pos_embd_global: self.pos_embd = PositionalEncodingFourier(dim=dims[0]) else: self.pos_embd = None self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers stem = nn.Sequential( nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), LayerNorm(dims[0], eps=1e-6, data_format="channels_first") ) self.downsample_layers.append(stem) for i in range(3): downsample_layer = nn.Sequential( LayerNorm(dims[i], eps=1e-6, data_format="channels_first"), nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), ) self.downsample_layers.append(downsample_layer) self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] cur = 0 for i in range(4): stage_blocks = [] for j in range(depths[i]): if j > depths[i] - global_block[i] - 1: if global_block_type[i] == 'SDTA': stage_blocks.append(SDTAEncoder(dim=dims[i], drop_path=dp_rates[cur + j], expan_ratio=expan_ratio, scales=d2_scales[i], use_pos_emb=use_pos_embd_xca[i], num_heads=heads[i])) else: raise NotImplementedError else: stage_blocks.append(ConvEncoder(dim=dims[i], drop_path=dp_rates[cur + j], layer_scale_init_value=layer_scale_init_value, expan_ratio=expan_ratio, kernel_size=kernel_sizes[i])) self.stages.append(nn.Sequential(*stage_blocks)) cur += depths[i] self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # Final norm layer self.head = nn.Linear(dims[-1], num_classes) self.apply(self._init_weights) self.head_dropout = nn.Dropout(class_drop_rate) self.head.weight.data.mul_(head_init_scale) self.head.bias.data.mul_(head_init_scale) def _init_weights(self, m): # TODO: MobileViT is using 'kaiming_normal' for initializing conv layers if isinstance(m, (nn.Conv2d, nn.Linear)): trunc_normal_(m.weight, std=.02) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, (LayerNorm, nn.LayerNorm)): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def forward_features(self, x): x = self.downsample_layers[0](x) x = self.stages[0](x) if self.pos_embd: B, C, H, W = x.shape x = x + self.pos_embd(B, H, W) for i in range(1, 4): x = self.downsample_layers[i](x) x = self.stages[i](x) return self.norm(x.mean([-2, -1])) # Global average pooling, (N, C, H, W) -> (N, C) def forward(self, x): x = self.forward_features(x) x = self.head(self.head_dropout(x)) return x class EdgeNeXtBNHS(nn.Module): def __init__(self, in_chans=3, num_classes=1000, depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], global_block=[0, 0, 0, 3], global_block_type=['None', 'None', 'None', 'SDTA_BN_HS'], drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., expan_ratio=4, kernel_sizes=[7, 7, 7, 7], heads=[8, 8, 8, 8], use_pos_embd_xca=[False, False, False, False], use_pos_embd_global=False, d2_scales=[2, 3, 4, 5], class_drop_rate=0.2, **kwargs): super().__init__() for g in global_block_type: assert g in ['None', 'SDTA_BN_HS'] if use_pos_embd_global: self.pos_embd = PositionalEncodingFourier(dim=dims[0]) else: self.pos_embd = None self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers stem = nn.Sequential( nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4, bias=False), nn.BatchNorm2d(dims[0]) ) self.downsample_layers.append(stem) for i in range(3): downsample_layer = nn.Sequential( nn.BatchNorm2d(dims[i]), nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2, bias=False), ) self.downsample_layers.append(downsample_layer) self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] cur = 0 for i in range(4): stage_blocks = [] for j in range(depths[i]): if j > depths[i] - global_block[i] - 1: if global_block_type[i] == 'SDTA_BN_HS': stage_blocks.append(SDTAEncoderBNHS(dim=dims[i], drop_path=dp_rates[cur + j], expan_ratio=expan_ratio, scales=d2_scales[i], use_pos_emb=use_pos_embd_xca[i], num_heads=heads[i])) else: raise NotImplementedError else: stage_blocks.append(ConvEncoderBNHS(dim=dims[i], drop_path=dp_rates[cur + j], layer_scale_init_value=layer_scale_init_value, expan_ratio=expan_ratio, kernel_size=kernel_sizes[i])) self.stages.append(nn.Sequential(*stage_blocks)) cur += depths[i] self.norm = nn.BatchNorm2d(dims[-1]) self.head = nn.Linear(dims[-1], num_classes) self.apply(self._init_weights) self.head_dropout = nn.Dropout(class_drop_rate) self.head.weight.data.mul_(head_init_scale) self.head.bias.data.mul_(head_init_scale) def _init_weights(self, m): # TODO: MobileViT is using 'kaiming_normal' for initializing conv layers if isinstance(m, (nn.Conv2d, nn.Linear)): trunc_normal_(m.weight, std=.02) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, (LayerNorm, nn.LayerNorm)): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def forward_features(self, x): x = self.downsample_layers[0](x) x = self.stages[0](x) if self.pos_embd: B, C, H, W = x.shape x = x + self.pos_embd(B, H, W) for i in range(1, 4): x = self.downsample_layers[i](x) x = self.stages[i](x) return self.norm(x).mean([-2, -1]) def forward(self, x): x = self.forward_features(x) x = self.head(self.head_dropout(x)) return x """ -- Main Models XX-Small -> 1.3M X-Small -> 2.3M Small -> 5.6M """ def edgenext_xx_small(pretrained=False, **kwargs): # 1.33M & 260.58M @ 256 resolution # 71.23% Top-1 accuracy # No AA, Color Jitter=0.4, No Mixup & Cutmix, DropPath=0.0, BS=4096, lr=0.006, multi-scale-sampler # Jetson FPS=51.66 versus 47.67 for MobileViT_XXS # For A100: FPS @ BS=1: 212.13 & @ BS=256: 7042.06 versus FPS @ BS=1: 96.68 & @ BS=256: 4624.71 for MobileViT_XXS model = EdgeNeXt(depths=[2, 2, 6, 2], dims=[24, 48, 88, 168], expan_ratio=4, global_block=[0, 1, 1, 1], global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'], use_pos_embd_xca=[False, True, False, False], kernel_sizes=[3, 5, 7, 9], heads=[4, 4, 4, 4], d2_scales=[2, 2, 3, 4], drop_path_rate=0.0, **kwargs) return model def edgenext_x_small(pretrained=False, **kwargs): # 2.34M & 538.0M @ 256 resolution # 75.00% Top-1 accuracy # No AA, No Mixup & Cutmix, DropPath=0.0, BS=4096, lr=0.006, multi-scale-sampler # Jetson FPS=31.61 versus 28.49 for MobileViT_XS # For A100: FPS @ BS=1: 179.55 & @ BS=256: 4404.95 versus FPS @ BS=1: 94.55 & @ BS=256: 2361.53 for MobileViT_XS model = EdgeNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 100, 192], expan_ratio=4, global_block=[0, 1, 1, 1], global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'], use_pos_embd_xca=[False, True, False, False], kernel_sizes=[3, 5, 7, 9], heads=[4, 4, 4, 4], d2_scales=[2, 2, 3, 4], **kwargs) return model def edgenext_small(pretrained=False, **kwargs): # 5.59M & 1260.59M @ 256 resolution # 79.43% Top-1 accuracy # AA=True, No Mixup & Cutmix, DropPath=0.1, BS=4096, lr=0.006, multi-scale-sampler # Jetson FPS=20.47 versus 18.86 for MobileViT_S # For A100: FPS @ BS=1: 172.33 & @ BS=256: 3010.25 versus FPS @ BS=1: 93.84 & @ BS=256: 1785.92 for MobileViT_S model = EdgeNeXt(depths=[3, 3, 9, 3], dims=[48, 96, 160, 304], expan_ratio=4, global_block=[0, 1, 1, 1], global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'], use_pos_embd_xca=[False, True, False, False], kernel_sizes=[3, 5, 7, 9], d2_scales=[2, 2, 3, 4], **kwargs) return model def edgenext_base(pretrained=False, **kwargs): # 18.51M & 3840.93M @ 256 resolution # 82.5% (normal) 83.7% (USI) Top-1 accuracy # AA=True, Mixup & Cutmix, DropPath=0.1, BS=4096, lr=0.006, multi-scale-sampler # Jetson FPS=xx.xx versus xx.xx for MobileViT_S # For A100: FPS @ BS=1: xxx.xx & @ BS=256: xxxx.xx model = EdgeNeXt(depths=[3, 3, 9, 3], dims=[80, 160, 288, 584], expan_ratio=4, global_block=[0, 1, 1, 1], global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'], use_pos_embd_xca=[False, True, False, False], kernel_sizes=[3, 5, 7, 9], d2_scales=[2, 2, 3, 4], drop_path_rate=0.1, **kwargs) return model """ Using BN & HSwish instead of LN & GeLU """ def edgenext_xx_small_bn_hs(pretrained=False, **kwargs): # 1.33M & 259.53M @ 256 resolution # 70.33% Top-1 accuracy # For A100: FPS @ BS=1: 219.66 & @ BS=256: 10359.98 model = EdgeNeXtBNHS(depths=[2, 2, 6, 2], dims=[24, 48, 88, 168], expan_ratio=4, global_block=[0, 1, 1, 1], global_block_type=['None', 'SDTA_BN_HS', 'SDTA_BN_HS', 'SDTA_BN_HS'], use_pos_embd_xca=[False, True, False, False], kernel_sizes=[3, 5, 7, 9], heads=[4, 4, 4, 4], d2_scales=[2, 2, 3, 4], **kwargs) return model def edgenext_x_small_bn_hs(pretrained=False, **kwargs): # 2.34M & 535.84M @ 256 resolution # 74.87% Top-1 accuracy # For A100: FPS @ BS=1: 179.25 & @ BS=256: 6059.59 model = EdgeNeXtBNHS(depths=[3, 3, 9, 3], dims=[32, 64, 100, 192], expan_ratio=4, global_block=[0, 1, 1, 1], global_block_type=['None', 'SDTA_BN_HS', 'SDTA_BN_HS', 'SDTA_BN_HS'], use_pos_embd_xca=[False, True, False, False], kernel_sizes=[3, 5, 7, 9], heads=[4, 4, 4, 4], d2_scales=[2, 2, 3, 4], drop_path_rate=0.2, **kwargs) return model def edgenext_small_bn_hs(pretrained=False, **kwargs): # 5.58M & 1257.28M @ 256 resolution # 78.39% Top-1 accuracy # For A100: FPS @ BS=1: 174.68 & @ BS=256: 3808.19 model = EdgeNeXtBNHS(depths=[3, 3, 9, 3], dims=[48, 96, 160, 304], expan_ratio=4, global_block=[0, 1, 1, 1], global_block_type=['None', 'SDTA_BN_HS', 'SDTA_BN_HS', 'SDTA_BN_HS'], use_pos_embd_xca=[False, True, False, False], kernel_sizes=[3, 5, 7, 9], d2_scales=[2, 2, 3, 4], drop_path_rate=0.2, **kwargs) return model