Skip to content

Commit

Permalink
Merge pull request #2277 from huggingface/more_tiny_test_models
Browse files Browse the repository at this point in the history
Adding some more tiny test models to train...
  • Loading branch information
rwightman authored Sep 22, 2024
2 parents 305a360 + 9067be6 commit e3242a5
Show file tree
Hide file tree
Showing 7 changed files with 140 additions and 12 deletions.
2 changes: 1 addition & 1 deletion timm/models/_efficientnet_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
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__)
Expand Down
2 changes: 1 addition & 1 deletion timm/models/byobnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -2355,7 +2355,7 @@ def _cfgr(url='', **kwargs):
'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),
),
})

Expand Down
37 changes: 37 additions & 0 deletions timm/models/convnext.py
Original file line number Diff line number Diff line change
Expand Up @@ -951,6 +951,20 @@ def _cfgv2(url='', **kwargs):
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),

})


Expand Down Expand Up @@ -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',
Expand Down
37 changes: 35 additions & 2 deletions timm/models/efficientnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,8 @@
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, \
Expand Down Expand Up @@ -1803,7 +1804,18 @@ def _cfg(url='', **kwargs):

"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),
})


Expand Down Expand Up @@ -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',
Expand Down
14 changes: 14 additions & 0 deletions timm/models/nfnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -607,6 +607,10 @@ def _dm_nfnet_cfg(
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'),
)


Expand Down Expand Up @@ -730,6 +734,11 @@ def _dcfg(url='', **kwargs):
'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)),
})


Expand Down Expand Up @@ -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)
28 changes: 22 additions & 6 deletions timm/models/resnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,8 @@
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
Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -490,8 +490,9 @@ def __init__(
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,
Expand All @@ -513,7 +514,7 @@ def __init__(
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)
Expand Down Expand Up @@ -1301,6 +1302,11 @@ def _gcfg(url='', **kwargs):
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'),
})


Expand Down Expand Up @@ -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',
Expand Down
32 changes: 30 additions & 2 deletions timm/models/vision_transformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -2014,7 +2014,13 @@ def _cfg(url: str = '', **kwargs) -> Dict[str, Any]:

'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 = [
Expand Down Expand Up @@ -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',
Expand Down

0 comments on commit e3242a5

Please sign in to comment.