add aa stem logic, minimal docstrings

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):