Tiny Test Models

I've recently trained a set of tiny test models (https://huggingface.co/collections/timm/timm-tiny-test-models-66f18bd70518277591a86cef) on ImageNet-1k covering several of the most popular architecture families.

It takes ~10 seconds to download all 13 pretrained weights and run one step of inference on each w/ a ratty old CPU (but fast internet connection). It will allow quick verification of model functionality, from pretrained weight download through every API feature of full size models. They differ from full size models in that they have lower default resolution and typically 1 block per stage, very narrow widths.

This is all well and good, but would anyone have any interest in these outside of tests? Well, this is where you come in. These are some of the smallest models that are decently trained on ImageNet-1k. They use a recent training recipe adapted from MobileNet-v4 (Conv-Small), a good recipe for squeezing accuracy from small models. The top-1 are by no means impressive, but the models do work well on fine-tune for small datasets, and I imagine they could work quite well for some reduced resource (embedded) applications or as part of reinforcement learning vision policies.

Let me know if you find any good applications for them outside of tests. Here's the summary of the model results, they were trained natively at 160x160, and most models see a small pickup at 192x192 by leveraging the train-test discrepancy.

ImageNet Accuracy

Througput @ 160x160 w/ torch.compile, mode='max-autotune', PyTorch 2.4.1, RTX4090

Througput @ 160x160 w/ torch.compile, mode='reduce-overhead', PyTorch 2.4.1, RTX4090

Througput @ 160x160 w/ torch.compile, mode='default', PyTorch 2.4.1, RTX4090

Output of python benchmark.py --amp --model 'test_*' --fast-norm --torchcompile:

Details

The model names above give some hint as to what they are, but I did explore some 'unique' architecture variations that are worth mentioning for any who might try them.

test_byobnet

A ByobNet (mix of EfficientNet / ResNet / DarkNet blocks)

• stage blocks = 1 * EdgeResidual (FusedMBConv), 1 * DarkBlock, 1 * ResNeXt Basic (group_size=32), 1 * ResNeXt Bottle (group_size=64)
• channels = 32, 64, 128, 256
• se_ratio = .25 (active in all blocks)
• act_layer = ReLU
• norm_layer = BatchNorm

test_convnext

A ConvNeXt

• stage depths = 1, 2, 4, 2
• channels = 24, 32, 48, 64
• DW kernel_size = 7, 7, 7, 7
• act_layer = GELU (tanh approx)
• norm_layer = LayerNorm

test_convnext2

A ConvNeXt

• stage depths = 1, 1, 1, 1
• channels = 32, 64, 96, 128
• DW kernel_size = 7, 7, 7, 7
• act_layer = GELU (tanh approx)
• norm_layer = LayerNorm

test_convnext3

A ConvNeXt w/ SiLU and varying kernel size

• stage depths = 1, 1, 1, 1
• channels = 32, 64, 96, 128
• DW kernel_size = 7, 5, 5, 3
• act_layer = SiLU
• norm_layer = LayerNorm

test_efficientnet

An EfficientNet w/ V2 block mix

• stage blocks = 1 * ConvBnAct, 2 * EdgeResidual (FusedMBConv), 2 * InvertedResidual (MBConv) w/ SE
• channles = 16, 24, 32, 48, 64
• kernel_size = 3x3 for all
• expansion = 4x for all
• stem_size = 24
• act_layer = SiLU
• norm_layer = BatchNorm

test_efficientnet_gn

An EfficientNet w/ V2 block mix and GroupNorm (group_size=8)

• See above but with norm_layer=GroupNorm

test_efficientnet_ln

An EfficientNet w/ V2 block mix and LayerNorm

• See above but with norm_layer=LayerNorm

test_efficientnet_evos

An EfficientNet w/ V2 block mix and EvoNorm-S

• See above but with EvoNormS for norm + act

test_nfnet

A NormFree Net:

• 4-stages, 1 block per stage
• channels = 32, 64, 96, 128
• group_size = 8
• bottle_ratio = 0.25
• se_ratio = 0.25
• act_layer = SiLU
• norm_layer = no norm, Scaled Weight Standardization is part of Convolution

test_resnet

A ResNet w/ mixed blocks:

• stage blocks = 1 * BasicBlock, 1 * BasicBlock, 1 * BottleNeck, 1 * BasicBlock
• channels = 32, 48, 48, 96
• deep 3x3 stem (aka ResNet-D)
• avg pool in downsample (aka ResNet-D)
• stem_width = 16
• act_layer = ReLU
• norm_layer = BatchNorm

test_vit

A vanilla ViT w/ class token:

• patch_size = 16
• embed_dim = 64
• num_heads = 2
• mlp_ratio = 3
• depth = 6
• act_layer = GELU
• norm_layer = LayerNorm

test_vit2

A ViT w/ global avg pool, 1 reg token, layer-scale (like timm SBB ViTs https://huggingface.co/collections/timm/searching-for-better-vit-baselines-663eb74f64f847d2f35a9c19):

• patch_size = 16
• embed_dim = 64
• num_heads = 2
• mlp_ratio = 3
• depth = 8
• act_layer = GELU
• norm_layer = LayerNorm

test_vit3

A ViT w/ attention-pool, 1 reg token, layer-scale.

• patch_size = 16
• embed_dim = 96
• num_heads = 3
• mlp_ratio = 2
• depth = 9
• act_layer = GELU
• norm_layer = LayerNorm