Add anti-aliasing support to mobilenetv3 and efficientnet family models. Update MobileNetV4 model defs, resolutions. Fix #599

* create_aa helper function centralized for all timm uses (resnet, convbnact helper)
* allow BlurPool w/ pre-defined channels (expand)
* mobilenetv4 UIB block using ConvNormAct layers for improved clarity, esp with AA added
* improve more mobilenetv3 and efficientnet related type annotations

Author: rwightman

timm/layers/__init__.py

@@ -4,7 +4,7 @@
 from .attention2d import MultiQueryAttention2d, Attention2d, MultiQueryAttentionV2
 from .attention_pool import AttentionPoolLatent
 from .attention_pool2d import AttentionPool2d, RotAttentionPool2d, RotaryEmbedding
-from .blur_pool import BlurPool2d
+from .blur_pool import BlurPool2d, create_aa
 from .classifier import ClassifierHead, create_classifier, NormMlpClassifierHead
 from .cond_conv2d import CondConv2d, get_condconv_initializer
 from .config import is_exportable, is_scriptable, is_no_jit, use_fused_attn, \

timm/layers/blur_pool.py

@@ -5,12 +5,16 @@
 
 Hacked together by Chris Ha and Ross Wightman
 """
+from functools import partial
+from typing import Optional, Type
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
+
 from .padding import get_padding
+from .typing import LayerType
 
 
 class BlurPool2d(nn.Module):
@@ -26,17 +30,62 @@ class BlurPool2d(nn.Module):
     Returns:
         torch.Tensor: the transformed tensor.
     """
-    def __init__(self, channels, filt_size=3, stride=2) -> None:
+    def __init__(
+            self,
+            channels: Optional[int] = None,
+            filt_size: int = 3,
+            stride: int = 2,
+            pad_mode: str = 'reflect',
+    ) -> None:
         super(BlurPool2d, self).__init__()
         assert filt_size > 1
         self.channels = channels
         self.filt_size = filt_size
         self.stride = stride
+        self.pad_mode = pad_mode
         self.padding = [get_padding(filt_size, stride, dilation=1)] * 4
+
         coeffs = torch.tensor((np.poly1d((0.5, 0.5)) ** (self.filt_size - 1)).coeffs.astype(np.float32))
-        blur_filter = (coeffs[:, None] * coeffs[None, :])[None, None, :, :].repeat(self.channels, 1, 1, 1)
+        blur_filter = (coeffs[:, None] * coeffs[None, :])[None, None, :, :]
+        if channels is not None:
+            blur_filter = blur_filter.repeat(self.channels, 1, 1, 1)
         self.register_buffer('filt', blur_filter, persistent=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.pad(x, self.padding, 'reflect')
-        return F.conv2d(x, self.filt, stride=self.stride, groups=self.channels)
+        x = F.pad(x, self.padding, mode=self.pad_mode)
+        if self.channels is None:
+            channels = x.shape[1]
+            weight = self.filt.expand(channels, 1, self.filt_size, self.filt_size)
+        else:
+            channels = self.channels
+            weight = self.filt
+        return F.conv2d(x, weight, stride=self.stride, groups=channels)
+
+
+def create_aa(
+        aa_layer: LayerType,
+        channels: Optional[int] = None,
+        stride: int = 2,
+        enable: bool = True,
+        noop: Optional[Type[nn.Module]] = nn.Identity
+) -> nn.Module:
+    """ Anti-aliasing """
+    if not aa_layer or not enable:
+        return noop() if noop is not None else None
+
+    if isinstance(aa_layer, str):
+        aa_layer = aa_layer.lower().replace('_', '').replace('-', '')
+        if aa_layer == 'avg' or aa_layer == 'avgpool':
+            aa_layer = nn.AvgPool2d
+        elif aa_layer == 'blur' or aa_layer == 'blurpool':
+            aa_layer = BlurPool2d
+        elif aa_layer == 'blurpc':
+            aa_layer = partial(BlurPool2d, pad_mode='constant')
+
+        else:
+            assert False, f"Unknown anti-aliasing layer ({aa_layer})."
+
+    try:
+        return aa_layer(channels=channels, stride=stride)
+    except TypeError as e:
+        return aa_layer(stride)

timm/layers/conv_bn_act.py

@@ -2,38 +2,51 @@
 
 Hacked together by / Copyright 2020 Ross Wightman
 """
-import functools
+from typing import Any, Dict, Optional, Type
+
 from torch import nn as nn
 
+from .typing import LayerType, PadType
+from .blur_pool import create_aa
 from .create_conv2d import create_conv2d
 from .create_norm_act import get_norm_act_layer
 
 
 class ConvNormAct(nn.Module):
     def __init__(
             self,
-            in_channels,
-            out_channels,
-            kernel_size=1,
-            stride=1,
-            padding='',
-            dilation=1,
-            groups=1,
-            bias=False,
-            apply_act=True,
-            norm_layer=nn.BatchNorm2d,
-            norm_kwargs=None,
-            act_layer=nn.ReLU,
-            act_kwargs=None,
-            drop_layer=None,
+            in_channels: int,
+            out_channels: int,
+            kernel_size: int = 1,
+            stride: int = 1,
+            padding: PadType = '',
+            dilation: int = 1,
+            groups: int = 1,
+            bias: bool = False,
+            apply_act: bool = True,
+            norm_layer: LayerType = nn.BatchNorm2d,
+            act_layer: LayerType = nn.ReLU,
+            drop_layer: Optional[Type[nn.Module]] = None,
+            conv_kwargs: Optional[Dict[str, Any]] = None,
+            norm_kwargs: Optional[Dict[str, Any]] = None,
+            act_kwargs: Optional[Dict[str, Any]] = None,
     ):
         super(ConvNormAct, self).__init__()
+        conv_kwargs = conv_kwargs or {}
         norm_kwargs = norm_kwargs or {}
         act_kwargs = act_kwargs or {}
 
         self.conv = create_conv2d(
-            in_channels, out_channels, kernel_size, stride=stride,
-            padding=padding, dilation=dilation, groups=groups, bias=bias)
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            **conv_kwargs,
+        )
 
         # NOTE for backwards compatibility with models that use separate norm and act layer definitions
         norm_act_layer = get_norm_act_layer(norm_layer, act_layer)
@@ -64,54 +77,53 @@ def forward(self, x):
 ConvBnAct = ConvNormAct
 
 
-def create_aa(aa_layer, channels, stride=2, enable=True):
-    if not aa_layer or not enable:
-        return nn.Identity()
-    if isinstance(aa_layer, functools.partial):
-        if issubclass(aa_layer.func, nn.AvgPool2d):
-            return aa_layer()
-        else:
-            return aa_layer(channels)
-    elif issubclass(aa_layer, nn.AvgPool2d):
-        return aa_layer(stride)
-    else:
-        return aa_layer(channels=channels, stride=stride)
-
-
 class ConvNormActAa(nn.Module):
     def __init__(
             self,
-            in_channels,
-            out_channels,
-            kernel_size=1,
-            stride=1,
-            padding='',
-            dilation=1,
-            groups=1,
-            bias=False,
-            apply_act=True,
-            norm_layer=nn.BatchNorm2d,
-            norm_kwargs=None,
-            act_layer=nn.ReLU,
-            act_kwargs=None,
-            aa_layer=None,
-            drop_layer=None,
+            in_channels: int,
+            out_channels: int,
+            kernel_size: int = 1,
+            stride: int = 1,
+            padding: PadType = '',
+            dilation: int = 1,
+            groups: int = 1,
+            bias: bool = False,
+            apply_act: bool = True,
+            norm_layer: LayerType = nn.BatchNorm2d,
+            act_layer: LayerType = nn.ReLU,
+            aa_layer: Optional[LayerType] = None,
+            drop_layer: Optional[Type[nn.Module]] = None,
+            conv_kwargs: Optional[Dict[str, Any]] = None,
+            norm_kwargs: Optional[Dict[str, Any]] = None,
+            act_kwargs: Optional[Dict[str, Any]] = None,
     ):
         super(ConvNormActAa, self).__init__()
         use_aa = aa_layer is not None and stride == 2
+        conv_kwargs = conv_kwargs or {}
         norm_kwargs = norm_kwargs or {}
         act_kwargs = act_kwargs or {}
 
         self.conv = create_conv2d(
-            in_channels, out_channels, kernel_size, stride=1 if use_aa else stride,
-            padding=padding, dilation=dilation, groups=groups, bias=bias)
+            in_channels, out_channels, kernel_size,
+            stride=1 if use_aa else stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            **conv_kwargs,
+        )
 
         # NOTE for backwards compatibility with models that use separate norm and act layer definitions
         norm_act_layer = get_norm_act_layer(norm_layer, act_layer)
         # NOTE for backwards (weight) compatibility, norm layer name remains `.bn`
         if drop_layer:
             norm_kwargs['drop_layer'] = drop_layer
-        self.bn = norm_act_layer(out_channels, apply_act=apply_act, act_kwargs=act_kwargs, **norm_kwargs)
+        self.bn = norm_act_layer(
+            out_channels,
+            apply_act=apply_act,
+            act_kwargs=act_kwargs,
+            **norm_kwargs,
+        )
         self.aa = create_aa(aa_layer, out_channels, stride=stride, enable=use_aa)
 
     @property