From 5078b28f8aec21bad91d0b866eb43bab33bc0bba Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Sun, 8 Jan 2023 18:17:17 -0800 Subject: [PATCH] More kwarg handling tweaks, maxvit_base_rw def added --- timm/models/densenet.py | 47 +++++++++++++++++++++--------- timm/models/dpn.py | 10 +++---- timm/models/maxxvit.py | 36 +++++++++++++++++++++++ timm/models/mobilevit.py | 13 +++++++-- timm/models/resnet.py | 4 +-- timm/models/resnetv2.py | 63 +++++++++++++++++++--------------------- 6 files changed, 117 insertions(+), 56 deletions(-) diff --git a/timm/models/densenet.py b/timm/models/densenet.py index e731f7b02a..ccbb491ca4 100644 --- a/timm/models/densenet.py +++ b/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" ` """ 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 diff --git a/timm/models/dpn.py b/timm/models/dpn.py index 6dbabad20f..29a7a7e8e6 100644 --- a/timm/models/dpn.py +++ b/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 diff --git a/timm/models/maxxvit.py b/timm/models/maxxvit.py index d68301b15c..7f9075b9a2 100644 --- a/timm/models/maxxvit.py +++ b/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) diff --git a/timm/models/mobilevit.py b/timm/models/mobilevit.py index 3d2ae84ab3..8e8f442851 100644 --- a/timm/models/mobilevit.py +++ b/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) diff --git a/timm/models/resnet.py b/timm/models/resnet.py index b3c041c436..200280b3ba 100644 --- a/timm/models/resnet.py +++ b/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 diff --git a/timm/models/resnetv2.py b/timm/models/resnetv2.py index d696b29177..41e29e1250 100644 --- a/timm/models/resnetv2.py +++ b/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))