Merge pull request #2277 from huggingface/more_tiny_test_models

Adding some more tiny test models to train...

timm/models/_efficientnet_builder.py

@@ -20,7 +20,7 @@ from timm.layers import CondConv2d, get_condconv_initializer, get_act_layer, get
 from ._efficientnet_blocks import *
 from ._manipulate import named_modules
 
-__all__ = ["EfficientNetBuilder", "decode_arch_def", "efficientnet_init_weights",
+__all__ = ["EfficientNetBuilder", "BlockArgs", "decode_arch_def", "efficientnet_init_weights",
            'resolve_bn_args', 'resolve_act_layer', 'round_channels', 'BN_MOMENTUM_TF_DEFAULT', 'BN_EPS_TF_DEFAULT']
 
 _logger = logging.getLogger(__name__)

timm/models/byobnet.py

@@ -2355,7 +2355,7 @@ default_cfgs = generate_default_cfgs({
     'test_byobnet.r160_in1k': _cfgr(
         hf_hub_id='timm/',
         first_conv='stem.conv',
-        input_size=(3, 160, 160), crop_pct=0.875, pool_size=(5, 5),
+        input_size=(3, 160, 160), crop_pct=0.95, pool_size=(5, 5),
     ),
 })
 

timm/models/convnext.py

@@ -951,6 +951,20 @@ default_cfgs = generate_default_cfgs({
         hf_hub_filename='open_clip_pytorch_model.bin',
         mean=OPENAI_CLIP_MEAN, std=OPENAI_CLIP_STD,
         input_size=(3, 256, 256), pool_size=(8, 8), crop_pct=1.0, num_classes=1024),
+
+    "test_convnext.r160_in1k": _cfg(
+        hf_hub_id='timm/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+        input_size=(3, 160, 160), pool_size=(5, 5), crop_pct=0.95),
+    "test_convnext2.r160_in1k": _cfg(
+        hf_hub_id='timm/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+        input_size=(3, 160, 160), pool_size=(5, 5), crop_pct=0.95),
+    "test_convnext3.r160_in1k": _cfg(
+        hf_hub_id='timm/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+        input_size=(3, 160, 160), pool_size=(5, 5), crop_pct=0.95),
+
 })
 
 
@@ -1146,6 +1160,29 @@ def convnextv2_huge(pretrained=False, **kwargs) -> ConvNeXt:
     return model
 
 
+@register_model
+def test_convnext(pretrained=False, **kwargs) -> ConvNeXt:
+    model_args = dict(depths=[1, 2, 4, 2], dims=[24, 32, 48, 64], norm_eps=kwargs.pop('norm_eps', 1e-5), act_layer='gelu_tanh')
+    model = _create_convnext('test_convnext', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def test_convnext2(pretrained=False, **kwargs) -> ConvNeXt:
+    model_args = dict(depths=[1, 1, 1, 1], dims=[32, 64, 96, 128], norm_eps=kwargs.pop('norm_eps', 1e-5), act_layer='gelu_tanh')
+    model = _create_convnext('test_convnext2', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def test_convnext3(pretrained=False, **kwargs) -> ConvNeXt:
+    model_args = dict(
+        depths=[1, 1, 1, 1], dims=[32, 64, 96, 128], norm_eps=kwargs.pop('norm_eps', 1e-5), kernel_sizes=(7, 5, 5, 3), act_layer='silu')
+    model = _create_convnext('test_convnext3', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+
 register_model_deprecations(__name__, {
     'convnext_tiny_in22ft1k': 'convnext_tiny.fb_in22k_ft_in1k',
     'convnext_small_in22ft1k': 'convnext_small.fb_in22k_ft_in1k',

timm/models/efficientnet.py

@@ -44,7 +44,8 @@ import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
-from timm.layers import create_conv2d, create_classifier, get_norm_act_layer, GroupNormAct, LayerType
+from timm.layers import create_conv2d, create_classifier, get_norm_act_layer, LayerType, \
+    GroupNormAct, LayerNormAct2d, EvoNorm2dS0
 from ._builder import build_model_with_cfg, pretrained_cfg_for_features
 from ._efficientnet_blocks import SqueezeExcite
 from ._efficientnet_builder import BlockArgs, EfficientNetBuilder, decode_arch_def, efficientnet_init_weights, \
@@ -1803,7 +1804,18 @@ default_cfgs = generate_default_cfgs({
 
     "test_efficientnet.r160_in1k": _cfg(
         hf_hub_id='timm/',
-        input_size=(3, 160, 160), pool_size=(5, 5)),
+        input_size=(3, 160, 160), pool_size=(5, 5), crop_pct=0.95),
+    "test_efficientnet_ln.r160_in1k": _cfg(
+        hf_hub_id='timm/',
+        input_size=(3, 160, 160), pool_size=(5, 5), crop_pct=0.95),
+    "test_efficientnet_gn.r160_in1k": _cfg(
+        hf_hub_id='timm/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+        input_size=(3, 160, 160), pool_size=(5, 5), crop_pct=0.95),
+    "test_efficientnet_evos.r160_in1k": _cfg(
+        hf_hub_id='timm/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+        input_size=(3, 160, 160), pool_size=(5, 5), crop_pct=0.95),
 })
 
 
@@ -2792,6 +2804,27 @@ def test_efficientnet(pretrained=False, **kwargs) -> EfficientNet:
     return model
 
 
+@register_model
+def test_efficientnet_gn(pretrained=False, **kwargs) -> EfficientNet:
+    model = _gen_test_efficientnet(
+        'test_efficientnet_gn', pretrained=pretrained, norm_layer=partial(GroupNormAct, group_size=8), **kwargs)
+    return model
+
+
+@register_model
+def test_efficientnet_ln(pretrained=False, **kwargs) -> EfficientNet:
+    model = _gen_test_efficientnet(
+        'test_efficientnet_ln', pretrained=pretrained, norm_layer=LayerNormAct2d, **kwargs)
+    return model
+
+
+@register_model
+def test_efficientnet_evos(pretrained=False, **kwargs) -> EfficientNet:
+    model = _gen_test_efficientnet(
+        'test_efficientnet_evos', pretrained=pretrained, norm_layer=partial(EvoNorm2dS0, group_size=8), **kwargs)
+    return model
+
+
 register_model_deprecations(__name__, {
     'tf_efficientnet_b0_ap': 'tf_efficientnet_b0.ap_in1k',
     'tf_efficientnet_b1_ap': 'tf_efficientnet_b1.ap_in1k',

timm/models/nfnet.py

@@ -607,6 +607,10 @@ model_cfgs = dict(
     nf_ecaresnet26=_nfres_cfg(depths=(2, 2, 2, 2), attn_layer='eca', attn_kwargs=dict()),
     nf_ecaresnet50=_nfres_cfg(depths=(3, 4, 6, 3), attn_layer='eca', attn_kwargs=dict()),
     nf_ecaresnet101=_nfres_cfg(depths=(3, 4, 23, 3), attn_layer='eca', attn_kwargs=dict()),
+
+    test_nfnet=_nfnet_cfg(
+        depths=(1, 1, 1, 1), channels=(32, 64, 96, 128), feat_mult=1.5, group_size=8, bottle_ratio=0.25,
+        attn_kwargs=dict(rd_ratio=0.25, rd_divisor=8), act_layer='silu'),
 )
 
 
@@ -730,6 +734,11 @@ default_cfgs = generate_default_cfgs({
     'nf_ecaresnet26': _dcfg(url='', first_conv='stem.conv'),
     'nf_ecaresnet50': _dcfg(url='', first_conv='stem.conv'),
     'nf_ecaresnet101': _dcfg(url='', first_conv='stem.conv'),
+
+    'test_nfnet.r160_in1k': _dcfg(
+        hf_hub_id='timm/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+        crop_pct=0.95, input_size=(3, 160, 160), pool_size=(5, 5)),
 })
 
 
@@ -1029,3 +1038,8 @@ def nf_ecaresnet101(pretrained=False, **kwargs) -> NormFreeNet:
     """ Normalization-Free ECA-ResNet101
     """
     return _create_normfreenet('nf_ecaresnet101', pretrained=pretrained, **kwargs)
+
+
+@register_model
+def test_nfnet(pretrained=False, **kwargs) -> NormFreeNet:
+    return _create_normfreenet('test_nfnet', pretrained=pretrained, **kwargs)

timm/models/resnet.py

@@ -16,8 +16,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
-from timm.layers import DropBlock2d, DropPath, AvgPool2dSame, BlurPool2d, GroupNorm, LayerType, create_attn, \
-    get_attn, get_act_layer, get_norm_layer, create_classifier, create_aa
+from timm.layers import DropBlock2d, DropPath, AvgPool2dSame, BlurPool2d, LayerType, create_attn, \
+    get_attn, get_act_layer, get_norm_layer, create_classifier, create_aa, to_ntuple
 from ._builder import build_model_with_cfg
 from ._features import feature_take_indices
 from ._manipulate import checkpoint_seq
@@ -286,7 +286,7 @@ def drop_blocks(drop_prob: float = 0.):
 
 
 def make_blocks(
-        block_fn: Union[BasicBlock, Bottleneck],
+        block_fns: Tuple[Union[BasicBlock, Bottleneck]],
         channels: Tuple[int, ...],
         block_repeats: Tuple[int, ...],
         inplanes: int,
@@ -304,7 +304,7 @@ def make_blocks(
     net_block_idx = 0
     net_stride = 4
     dilation = prev_dilation = 1
-    for stage_idx, (planes, num_blocks, db) in enumerate(zip(channels, block_repeats, drop_blocks(drop_block_rate))):
+    for stage_idx, (block_fn, planes, num_blocks, db) in enumerate(zip(block_fns, channels, block_repeats, drop_blocks(drop_block_rate))):
         stage_name = f'layer{stage_idx + 1}'  # never liked this name, but weight compat requires it
         stride = 1 if stage_idx == 0 else 2
         if net_stride >= output_stride:
@@ -490,8 +490,9 @@ class ResNet(nn.Module):
                 self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
 
         # Feature Blocks
+        block_fns = to_ntuple(len(channels))(block)
         stage_modules, stage_feature_info = make_blocks(
-            block,
+            block_fns,
             channels,
             layers,
             inplanes,
@@ -513,7 +514,7 @@ class ResNet(nn.Module):
         self.feature_info.extend(stage_feature_info)
 
         # Head (Pooling and Classifier)
-        self.num_features = self.head_hidden_size = channels[-1] * block.expansion
+        self.num_features = self.head_hidden_size = channels[-1] * block_fns[-1].expansion
         self.global_pool, self.fc = create_classifier(self.num_features, self.num_classes, pool_type=global_pool)
 
         self.init_weights(zero_init_last=zero_init_last)
@@ -1301,6 +1302,11 @@ default_cfgs = generate_default_cfgs({
         hf_hub_id='timm/',
         url='https://github.com/rwightman/pytorch-pretrained-gluonresnet/releases/download/v0.1/gluon_senet154-70a1a3c0.pth',
         first_conv='conv1.0'),
+
+    'test_resnet.r160_in1k': _cfg(
+        hf_hub_id='timm/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), crop_pct=0.95,
+        input_size=(3, 160, 160), pool_size=(5, 5), first_conv='conv1.0'),
 })
 
 
@@ -2040,6 +2046,16 @@ def resnetrs420(pretrained: bool = False, **kwargs) -> ResNet:
     return _create_resnet('resnetrs420', pretrained, **dict(model_args, **kwargs))
 
 
+@register_model
+def test_resnet(pretrained: bool = False, **kwargs) -> ResNet:
+    """Constructs a tiny ResNet test model.
+    """
+    model_args = dict(
+        block=[BasicBlock, BasicBlock, Bottleneck, BasicBlock], layers=(1, 1, 1, 1),
+        stem_width=16, stem_type='deep', avg_down=True, channels=(32, 48, 48, 96))
+    return _create_resnet('test_resnet', pretrained, **dict(model_args, **kwargs))
+
+
 register_model_deprecations(__name__, {
     'tv_resnet34': 'resnet34.tv_in1k',
     'tv_resnet50': 'resnet50.tv_in1k',

timm/models/vision_transformer.py

@@ -2014,7 +2014,13 @@ default_cfgs = {
 
     'test_vit.r160_in1k': _cfg(
         hf_hub_id='timm/',
-        input_size=(3, 160, 160), crop_pct=0.875),
+        input_size=(3, 160, 160), crop_pct=0.95),
+    'test_vit2.r160_in1k': _cfg(
+        hf_hub_id='timm/',
+        input_size=(3, 160, 160), crop_pct=0.95),
+    'test_vit3.r160_in1k': _cfg(
+        #hf_hub_id='timm/',
+        input_size=(3, 160, 160), crop_pct=0.95),
 }
 
 _quick_gelu_cfgs = [
@@ -3211,11 +3217,33 @@ def vit_so150m_patch16_reg4_gap_256(pretrained: bool = False, **kwargs) -> Visio
 def test_vit(pretrained: bool = False, **kwargs) -> VisionTransformer:
     """ ViT Test
     """
-    model_args = dict(patch_size=16, embed_dim=64, depth=6, num_heads=2, mlp_ratio=3)
+    model_args = dict(patch_size=16, embed_dim=64, depth=6, num_heads=2, mlp_ratio=3, dynamic_img_size=True)
     model = _create_vision_transformer('test_vit', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
+@register_model
+def test_vit2(pretrained: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT Test
+    """
+    model_args = dict(
+        patch_size=16, embed_dim=64, depth=8, num_heads=2, mlp_ratio=3,
+        class_token=False, reg_tokens=1, global_pool='avg', init_values=1e-5, dynamic_img_size=True)
+    model = _create_vision_transformer('test_vit2', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def test_vit3(pretrained: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT Test
+    """
+    model_args = dict(
+        patch_size=16, embed_dim=96, depth=10, num_heads=3, mlp_ratio=2,
+        class_token=False, reg_tokens=1, global_pool='map', init_values=1e-5)
+    model = _create_vision_transformer('test_vit3', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
 register_model_deprecations(__name__, {
     'vit_tiny_patch16_224_in21k': 'vit_tiny_patch16_224.augreg_in21k',
     'vit_small_patch32_224_in21k': 'vit_small_patch32_224.augreg_in21k',