More kwarg handling tweaks, maxvit_base_rw def added

pull/1633/head
Ross Wightman 2 years ago committed by Ross Wightman
parent c0d7388a1b
commit 5078b28f8a

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

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

@ -1680,6 +1680,26 @@ model_cfgs = dict(
init_values=1e-6, 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( maxvit_tiny_pm_256=MaxxVitCfg(
embed_dim=(64, 128, 256, 512), embed_dim=(64, 128, 256, 512),
@ -1862,6 +1882,12 @@ default_cfgs = generate_default_cfgs({
'maxvit_rmlp_small_rw_256': _cfg( 'maxvit_rmlp_small_rw_256': _cfg(
url='', url='',
input_size=(3, 256, 256), pool_size=(8, 8)), 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)), '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) 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 @register_model
def maxvit_tiny_pm_256(pretrained=False, **kwargs): def maxvit_tiny_pm_256(pretrained=False, **kwargs):
return _create_maxxvit('maxvit_tiny_pm_256', pretrained=pretrained, **kwargs) return _create_maxxvit('maxvit_tiny_pm_256', pretrained=pretrained, **kwargs)

@ -266,9 +266,16 @@ class MobileVitBlock(nn.Module):
self.transformer = nn.Sequential(*[ self.transformer = nn.Sequential(*[
TransformerBlock( TransformerBlock(
transformer_dim, mlp_ratio=mlp_ratio, num_heads=num_heads, qkv_bias=True, transformer_dim,
attn_drop=attn_drop, drop=drop, drop_path=drop_path_rate, mlp_ratio=mlp_ratio,
act_layer=layers.act, norm_layer=transformer_norm_layer) 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) for _ in range(transformer_depth)
]) ])
self.norm = transformer_norm_layer(transformer_dim) self.norm = transformer_norm_layer(transformer_dim)

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

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

Loading…
Cancel
Save