More kwarg handling tweaks, maxvit_base_rw def added
parent
c0d7388a1b
commit
5078b28f8a
|
@ -12,7 +12,7 @@ import torch.utils.checkpoint as cp
|
|||
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 @@ class DenseNet(nn.Module):
|
|||
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
|
||||
|
||||
|
||||
|
|
|
@ -178,21 +178,21 @@ class DualPathBlock(nn.Module):
|
|||
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
|
||||
|
|
|
@ -1680,6 +1680,26 @@ model_cfgs = dict(
|
|||
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 @@ default_cfgs = generate_default_cfgs({
|
|||
'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)
|
||||
|
|
|
@ -266,9 +266,16 @@ class MobileVitBlock(nn.Module):
|
|||
|
||||
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)
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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))
|
||||
|
|
Loading…
Reference in New Issue