More kwarg handling tweaks, maxvit_base_rw def added

timm/models/densenet.py

@@ -12,7 +12,7 @@
 from torch.jit.annotations import List
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
-from timm.layers import BatchNormAct2d, create_norm_act_layer, BlurPool2d, create_classifier
+from timm.layers import BatchNormAct2d, get_norm_act_layer, BlurPool2d, create_classifier
 from ._builder import build_model_with_cfg
 from ._manipulate import MATCH_PREV_GROUP
 from ._registry import register_model
@@ -115,8 +115,15 @@ class DenseBlock(nn.ModuleDict):
     _version = 2
 
     def __init__(
-            self, num_layers, num_input_features, bn_size, growth_rate, norm_layer=BatchNormAct2d,
-            drop_rate=0., memory_efficient=False):
+            self,
+            num_layers,
+            num_input_features,
+            bn_size,
+            growth_rate,
+            norm_layer=BatchNormAct2d,
+            drop_rate=0.,
+            memory_efficient=False,
+    ):
         super(DenseBlock, self).__init__()
         for i in range(num_layers):
             layer = DenseLayer(
@@ -165,12 +172,25 @@ class DenseNet(nn.Module):
     """
 
     def __init__(
-            self, growth_rate=32, block_config=(6, 12, 24, 16), num_classes=1000, in_chans=3, global_pool='avg',
-            bn_size=4, stem_type='', norm_layer=BatchNormAct2d, aa_layer=None, drop_rate=0,
-            memory_efficient=False, aa_stem_only=True):
+            self,
+            growth_rate=32,
+            block_config=(6, 12, 24, 16),
+            num_classes=1000,
+            in_chans=3,
+            global_pool='avg',
+            bn_size=4,
+            stem_type='',
+            act_layer='relu',
+            norm_layer='batchnorm2d',
+            aa_layer=None,
+            drop_rate=0,
+            memory_efficient=False,
+            aa_stem_only=True,
+    ):
         self.num_classes = num_classes
         self.drop_rate = drop_rate
         super(DenseNet, self).__init__()
+        norm_layer = get_norm_act_layer(norm_layer, act_layer=act_layer)
 
         # Stem
         deep_stem = 'deep' in stem_type  # 3x3 deep stem
@@ -226,8 +246,11 @@ def __init__(
                     dict(num_chs=num_features, reduction=current_stride, module='features.' + module_name)]
                 current_stride *= 2
                 trans = DenseTransition(
-                    num_input_features=num_features, num_output_features=num_features // 2,
-                    norm_layer=norm_layer, aa_layer=transition_aa_layer)
+                    num_input_features=num_features,
+                    num_output_features=num_features // 2,
+                    norm_layer=norm_layer,
+                    aa_layer=transition_aa_layer,
+                )
                 self.features.add_module(f'transition{i + 1}', trans)
                 num_features = num_features // 2
 
@@ -322,8 +345,8 @@ def densenetblur121d(pretrained=False, **kwargs):
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
     """
     model = _create_densenet(
-        'densenetblur121d', growth_rate=32, block_config=(6, 12, 24, 16), pretrained=pretrained, stem_type='deep',
-        aa_layer=BlurPool2d, **kwargs)
+        'densenetblur121d', growth_rate=32, block_config=(6, 12, 24, 16), pretrained=pretrained,
+        stem_type='deep', aa_layer=BlurPool2d, **kwargs)
     return model
 
 
@@ -382,11 +405,9 @@ def densenet264(pretrained=False, **kwargs):
 def densenet264d_iabn(pretrained=False, **kwargs):
     r"""Densenet-264 model with deep stem and Inplace-ABN
     """
-    def norm_act_fn(num_features, **kwargs):
-        return create_norm_act_layer('iabn', num_features, act_layer='leaky_relu', **kwargs)
     model = _create_densenet(
         'densenet264d_iabn', growth_rate=48, block_config=(6, 12, 64, 48), stem_type='deep',
-        norm_layer=norm_act_fn, pretrained=pretrained, **kwargs)
+        norm_layer='iabn', act_layer='leaky_relu', pretrained=pretrained, **kwargs)
     return model
 
 

timm/models/dpn.py

@@ -178,21 +178,21 @@ def forward(self, x) -> Tuple[torch.Tensor, torch.Tensor]:
 class DPN(nn.Module):
     def __init__(
             self,
-            num_classes=1000,
-            in_chans=3,
-            output_stride=32,
-            global_pool='avg',
             k_sec=(3, 4, 20, 3),
             inc_sec=(16, 32, 24, 128),
             k_r=96,
             groups=32,
+            num_classes=1000,
+            in_chans=3,
+            output_stride=32,
+            global_pool='avg',
             small=False,
             num_init_features=64,
             b=False,
             drop_rate=0.,
             norm_layer='batchnorm2d',
             act_layer='relu',
-            fc_act_layer=nn.ELU,
+            fc_act_layer='elu',
     ):
         super(DPN, self).__init__()
         self.num_classes = num_classes

timm/models/maxxvit.py

@@ -1680,6 +1680,26 @@ def _tf_cfg():
             init_values=1e-6,
         ),
     ),
+    maxvit_rmlp_base_rw_224=MaxxVitCfg(
+        embed_dim=(96, 192, 384, 768),
+        depths=(2, 6, 14, 2),
+        block_type=('M',) * 4,
+        stem_width=(32, 64),
+        head_hidden_size=768,
+        **_rw_max_cfg(
+            rel_pos_type='mlp',
+        ),
+    ),
+    maxvit_rmlp_base_rw_384=MaxxVitCfg(
+        embed_dim=(96, 192, 384, 768),
+        depths=(2, 6, 14, 2),
+        block_type=('M',) * 4,
+        stem_width=(32, 64),
+        head_hidden_size=768,
+        **_rw_max_cfg(
+            rel_pos_type='mlp',
+        ),
+    ),
 
     maxvit_tiny_pm_256=MaxxVitCfg(
         embed_dim=(64, 128, 256, 512),
@@ -1862,6 +1882,12 @@ def _cfg(url='', **kwargs):
     'maxvit_rmlp_small_rw_256': _cfg(
         url='',
         input_size=(3, 256, 256), pool_size=(8, 8)),
+    'maxvit_rmlp_base_rw_2244': _cfg(
+        url='',
+    ),
+    'maxvit_rmlp_base_rw_384': _cfg(
+        url='',
+        input_size=(3, 384, 384), pool_size=(12, 12)),
 
     'maxvit_tiny_pm_256': _cfg(url='', input_size=(3, 256, 256), pool_size=(8, 8)),
 
@@ -2091,6 +2117,16 @@ def maxvit_rmlp_small_rw_256(pretrained=False, **kwargs):
     return _create_maxxvit('maxvit_rmlp_small_rw_256', pretrained=pretrained, **kwargs)
 
 
+@register_model
+def maxvit_rmlp_base_rw_224(pretrained=False, **kwargs):
+    return _create_maxxvit('maxvit_rmlp_base_rw_224', pretrained=pretrained, **kwargs)
+
+
+@register_model
+def maxvit_rmlp_base_rw_384(pretrained=False, **kwargs):
+    return _create_maxxvit('maxvit_rmlp_base_rw_384', pretrained=pretrained, **kwargs)
+
+
 @register_model
 def maxvit_tiny_pm_256(pretrained=False, **kwargs):
     return _create_maxxvit('maxvit_tiny_pm_256', pretrained=pretrained, **kwargs)

timm/models/mobilevit.py

@@ -266,9 +266,16 @@ def __init__(
 
         self.transformer = nn.Sequential(*[
             TransformerBlock(
-                transformer_dim, mlp_ratio=mlp_ratio, num_heads=num_heads, qkv_bias=True,
-                attn_drop=attn_drop, drop=drop, drop_path=drop_path_rate,
-                act_layer=layers.act, norm_layer=transformer_norm_layer)
+                transformer_dim,
+                mlp_ratio=mlp_ratio,
+                num_heads=num_heads,
+                qkv_bias=True,
+                attn_drop=attn_drop,
+                drop=drop,
+                drop_path=drop_path_rate,
+                act_layer=layers.act,
+                norm_layer=transformer_norm_layer,
+            )
             for _ in range(transformer_depth)
         ])
         self.norm = transformer_norm_layer(transformer_dim)

timm/models/resnet.py

@@ -1298,7 +1298,7 @@ def ecaresnet50d_pruned(pretrained=False, **kwargs):
     model_args = dict(
         block=Bottleneck, layers=[3, 4, 6, 3], stem_width=32, stem_type='deep', avg_down=True,
         block_args=dict(attn_layer='eca'))
-    return _create_resnet('ecaresnet50d_pruned', pretrained, pruned=True, **model_args)
+    return _create_resnet('ecaresnet50d_pruned', pretrained, pruned=True, **dict(model_args, **kwargs))
 
 
 @register_model
@@ -1340,7 +1340,7 @@ def ecaresnet101d_pruned(pretrained=False, **kwargs):
     model_args = dict(
         block=Bottleneck, layers=[3, 4, 23, 3], stem_width=32, stem_type='deep', avg_down=True,
         block_args=dict(attn_layer='eca'))
-    return _create_resnet('ecaresnet101d_pruned', pretrained, pruned=True, **model_args)
+    return _create_resnet('ecaresnet101d_pruned', pretrained, pruned=True, **dict(model_args, **kwargs))
 
 
 @register_model

timm/models/resnetv2.py

@@ -746,86 +746,83 @@ def resnetv2_152x2_bit_teacher_384(pretrained=False, **kwargs):
 
 @register_model
 def resnetv2_50(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_50', pretrained=pretrained,
-        layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d, **kwargs)
+    model_args = dict(layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d)
+    return _create_resnetv2('resnetv2_50', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_50d(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_50d', pretrained=pretrained,
+    model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d,
-        stem_type='deep', avg_down=True, **kwargs)
+        stem_type='deep', avg_down=True)
+    return _create_resnetv2('resnetv2_50d', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_50t(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_50t', pretrained=pretrained,
+    model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d,
-        stem_type='tiered', avg_down=True, **kwargs)
+        stem_type='tiered', avg_down=True)
+    return _create_resnetv2('resnetv2_50t', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_101(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_101', pretrained=pretrained,
-        layers=[3, 4, 23, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d, **kwargs)
+    model_args = dict(layers=[3, 4, 23, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d)
+    return _create_resnetv2('resnetv2_101', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_101d(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_101d', pretrained=pretrained,
+    model_args = dict(
         layers=[3, 4, 23, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d,
-        stem_type='deep', avg_down=True, **kwargs)
+        stem_type='deep', avg_down=True)
+    return _create_resnetv2('resnetv2_101d', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_152(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_152', pretrained=pretrained,
-        layers=[3, 8, 36, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d, **kwargs)
+    model_args = dict(layers=[3, 8, 36, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d)
+    return _create_resnetv2('resnetv2_152', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_152d(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_152d', pretrained=pretrained,
+    model_args = dict(
         layers=[3, 8, 36, 3], conv_layer=create_conv2d, norm_layer=BatchNormAct2d,
-        stem_type='deep', avg_down=True, **kwargs)
+        stem_type='deep', avg_down=True)
+    return _create_resnetv2('resnetv2_152d', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 # Experimental configs (may change / be removed)
 
 @register_model
 def resnetv2_50d_gn(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_50d_gn', pretrained=pretrained,
+    model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=GroupNormAct,
-        stem_type='deep', avg_down=True, **kwargs)
+        stem_type='deep', avg_down=True)
+    return _create_resnetv2('resnetv2_50d_gn', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_50d_evob(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_50d_evob', pretrained=pretrained,
+    model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=EvoNorm2dB0,
-        stem_type='deep', avg_down=True, zero_init_last=True, **kwargs)
+        stem_type='deep', avg_down=True, zero_init_last=True)
+    return _create_resnetv2('resnetv2_50d_evob', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_50d_evos(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_50d_evos', pretrained=pretrained,
+    model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=EvoNorm2dS0,
-        stem_type='deep', avg_down=True, **kwargs)
+        stem_type='deep', avg_down=True)
+    return _create_resnetv2('resnetv2_50d_evos', pretrained=pretrained, **dict(model_args, **kwargs))
 
 
 @register_model
 def resnetv2_50d_frn(pretrained=False, **kwargs):
-    return _create_resnetv2(
-        'resnetv2_50d_frn', pretrained=pretrained,
+    model_args = dict(
         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=FilterResponseNormTlu2d,
-        stem_type='deep', avg_down=True, **kwargs)
+        stem_type='deep', avg_down=True)
+    return _create_resnetv2('resnetv2_50d_frn', pretrained=pretrained, **dict(model_args, **kwargs))