New Model Architectures - Implementation and Documentation Details

🚀 The feature

When adding a new model architecture there are some design/implementation details and documentation requirements that need to be taken into account. This issue intents to track such details in a dynamic manner, as it is possible to change over time.

Motivation, pitch

New Model Architectures - Implementation Details

Model development and training steps

When developing a new model there are some details not to be missed:

• Implement a model factory function for each of the model variants

• in the module constructor, pass layer constructor instead of instance for configurable layers like norm, activation, and log the api usage with _log_api_usage_once(self)

• fuse layers together with existing common blocks if possible; For example consecutive conv, bn, activation layers could be replaced by ConvNormActivation

• define __all__ in the beginning of the model file to expose model factory functions; import model public APIs (e.g. factory methods) in torchvision/models/__init__.py

• create the model builder using the new API and add it to the prototype area. Here is an example on how to do this. The new API requires adding more information about the weights such as the preprocessing transforms necessary for using the model, meta-data about the model, etc

• Make sure you write tests for the model itself (see _check_input_backprop, _model_params and _model_params in test/test_models.py) and for any new operators/transforms or important functions that you introduce

• the new model should be torch scriptable (using torch.jit.script)

• the new model should be fx compatible (using torch.fx.symbolic_trace)

Note that this list is not exhaustive and there are details here related to the code quality etc, but these are rules that apply in all PRs (see Contributing to TorchVision).

Once the model is implemented, you need to train the model using the reference scripts. For example, in order to train a classification resnet18 model you would:

1. go to references/classification

2. run the train command (for example torchrun --nproc_per_node=8 train.py --model resnet18)

After training the model, select the best checkpoint and estimate its accuracy with a batch size of 1 on a single GPU. This helps us get better measurements about the accuracy of the models and avoid variants introduced due to batch padding (read here for more details).

Finally, run the model test to generate expected model files for testing. Please include those generated files in the PR as well.:

EXPECTTEST_ACCEPT=1 pytest test/test_models.py -k {model_name}

Documentation and Pytorch Hub

• docs/source/models.rst:

  • add the model to the corresponding section (classification/detection/video etc.)

  • describe how to construct the model variants (with and without pre-trained weights)

  • add model metrics and reference to the original paper

• hubconf.py:

  • import the model factory functions

  • submit a PR to https://github.com/pytorch/hub with a model page (or update an existing one)

• README.md under the reference script folder:

  • command(s) to train the model

Alternatives

No response

Additional context

No response

Start by reading test/test_models.py, docs/source/models.rst, hubconf.py, and the README under references/classification. The issue also points to the model factory, prototype API, training command, and generated model-test files as relevant entry points. Done would require implementing and testing a new architecture, training and measuring it, generating expected files, and documenting it across the named locations.