is there some template for creating my own model ?

[aliutkus]

Hi there,

thanks for the amazing work

before posting any feature request, I wanted to open a discussion, because I may be just missing stuff.

I was wondering:

• is there some specific section in the doc that is dedicated to creating my own new model to work on timm ?
• is there some template of a minimal-working example of an empty model to start with ? That would be super practical

best

[rwightman]

@aliutkus that's a good question and I should have a guide for that, right now, the only guide is the existing models and the unit tests, the unit tests cover the core model API so if you created a model in timm and it's implemented correctly, it'll get picked up by the unit tests and pass.

The models are fairly standard PyTorch models (not much different from torchvision) with some extras and some uniformity for certain args.

To add a model in timm you need:

• model class that inherits from nn.Module (ie ResNet) and implements the architecture
• a model 'entrypoint function' w/ a specific architecture definition (ie resnet50) that instantiates the class w/ args for a variant of that architecture. All entrypoints should be decorated with @register_model so that they are added to the timm registry and the create fn etc works.
• a pretrained cfg entry (in currently called default_cfgs dictionary in impl file) that defines weight urls, and details of the dataset / classifier head relevant for the weights, num_classes, dataset mean/std stats, etc

For init args, all models should have these present:
num_classes=1000, in_chans=3, drop_rate=0.0

Also common for many (but not all) are drop_path_rate=0., global_pool='avg', output_stride=32

Models that have customizable activations and norm layers usually have act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d

Model should implement common functionality to create the model or reset/get the classifier head after creation.

Most of my though re interface was when all models were convnets, now that a lot are vision transformers or mlp models, there are two classes of models. I currently have some work underway to make the vit-like models a bit closer to the CNN on a few of the API details around classifier / pooling / feature extraction.

There are a bunch of non-obvious details re setting up the feature extraction interface correctly, that's non-trivial to describe....

[rwightman]

In past year or so I unified the model building, so each architecture has an internal _build function that calls build_model_with_cfg( ... this takes care of loading pretrained weights (possibly from different sources), resolving some args from pretrained cfg vs passed in via factory or entrypoint kwargs, setting up feature extraction wrapper for most models (some need custom code for this still)...

[rwightman]

Within the loose constraints by some common __init__ args, foward_features, get_classifier, reset_classifier there is quite a lot of room for different model building approaches.

• resnet.py was adapted from origina torchvision ResNet impl, but added support for more regularization, different stem configs, shortcut downsample, combined ResNeXt into it, etc.
• efficientnet.py(and mobilenetv3) includes its own set of block definitions, and a builder that puts together stages based on string defs and covers quite an variety of mbconv and depthwise_separable based models
• byobnet.py is a more recent model family I put together with a number of block options and dataclass based configs to define quite a few different types of models... I lean towards some of these ideas now. nfnet.py had some similar config ideas. One can implement ResNets, RegNet, ResNeXt and many other ResNet like models with a few extra bells and whistles with the byob and byoa (adding self-attn blocks as an option) building blocks.
• vision_transformer.py was my first vit type model, and the template for quite a few 3rd party PyTorch impl of later model variants. It has some extra noise for supporting embedding upscale, loading the JAX .npy weights, etc. One thing I want to fix here is move the token feature vs avg pool across the 1d sequence into forward() instead of in forward_features() so fwd features always returns a 1d sequence or 2d feature map and forward always does the final global pool/reduction op and classifier.

[aliutkus]

Thanks for the very detailed and great answers. I mark this as an answer, but continue the discussion

I must admit that it was indeed ok to start writing my own model rapidly: the overhead due to the framework is minimal. I managed to do that quite easily.

The daunting part precisely is about the factory: model creation, parameters handling, the create_model part, etc.

As a suggestion, it could be great to have some dummy.py file where everything is ready for a stupidly simple model that just takes the input and produces some prediction, after, say, some 1 layer CNN or something like that, that could easily be replaced by the practitioner.
This file could comprise all the registering, the configuration stuff etc (when you start with a new model, there's no pretrained weights online, etc). This would be wonderful to start with

Also, this would make it easier for anyone to propose PR to augment timm with his own model !

thanks again for the great work and the time you took to answer my question