diff --git a/timm/models/mobilenetv3.py b/timm/models/mobilenetv3.py index 026ce363..dcccf4ce 100644 --- a/timm/models/mobilenetv3.py +++ b/timm/models/mobilenetv3.py @@ -41,6 +41,7 @@ default_cfgs = { url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mobilenetv3_large_100_ra-f55367f5.pth'), 'mobilenetv3_large_075_aa': _cfg(url=''), 'mobilenetv3_large_100_aa': _cfg(url='https://storage.googleapis.com/cinemanet-models/pretrained/mobilenetv3_large_100_aa_224x224_ema.pth'), + 'mobilenetv3_large_100_aa_stem': _cfg(url=''), 'mobilenetv3_large_100_miil': _cfg( interpolation='bilinear', mean=(0, 0, 0), std=(1, 1, 1), url='https://miil-public-eu.oss-eu-central-1.aliyuncs.com/model-zoo/ImageNet_21K_P/models/timm/mobilenetv3_large_100_1k_miil_78_0.pth'), @@ -87,7 +88,7 @@ class MobileNetV3(nn.Module): def __init__(self, block_args, num_classes=1000, in_chans=3, stem_size=16, num_features=1280, head_bias=True, channel_multiplier=1.0, pad_type='', act_layer=nn.ReLU, drop_rate=0., drop_path_rate=0., - se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None, global_pool='avg', aa_layer=None): + se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None, global_pool='avg', aa_layer=None, aa_stem=None): super(MobileNetV3, self).__init__() self.num_classes = num_classes @@ -97,6 +98,7 @@ class MobileNetV3(nn.Module): # Stem stem_size = round_channels(stem_size, channel_multiplier) self.conv_stem = create_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type) + self.conv_stem_aa = aa_stem(in_chans) if aa_stem else None self.bn1 = norm_layer(stem_size, **norm_kwargs) self.act1 = act_layer(inplace=True) @@ -135,6 +137,8 @@ class MobileNetV3(nn.Module): def forward_features(self, x): x = self.conv_stem(x) + if self.conv_stem_aa is not None: + x = self.conv_stem_aa(x) x = self.bn1(x) x = self.act1(x) x = self.blocks(x) @@ -416,10 +420,17 @@ def mobilenetv3_large_075_aa(pretrained=False, aa_layer=BlurPool2d, **kwargs): @register_model def mobilenetv3_large_100_aa(pretrained=False, aa_layer=BlurPool2d, **kwargs): - """ MobileNet V3 """ + """ MobileNet V3 w/ Blur Pooling of IR Blocks """ model = _gen_mobilenet_v3('mobilenetv3_large_100_aa', 1.0, pretrained=pretrained, aa_layer=aa_layer, **kwargs) return model +@register_model +def mobilenetv3_large_100_aa_stem(pretrained=False, aa_layer=BlurPool2d, aa_stem=BlurPool2d, **kwargs): + """ MobileNet V3 w/ Blur Pooling of IR Blocks & Conv Stem """ + model = _gen_mobilenet_v3('mobilenetv3_large_100_aa_stem', 1.0, pretrained=pretrained, + aa_layer=aa_layer, aa_stem=aa_stem, **kwargs) + return model + @register_model def mobilenetv3_rw(pretrained=False, **kwargs):