TIMM (pyTorch IMage Models) is a useful collection of implemented and trained backbones. Standalone they can be used for image classification which is also the intended use case of any of the pre-trained models. Despite various advances in many subfields of deep learning based computer vision, the plain classical application of “image classification” is still very useful and widely applied. This has many reasons, some are performance related, others are that segmentation or object detection is not always necessary if e.g. anomalies are to be detected.

Deploying models trained on custom datasets is a tiny bit less straight forward unless ONNX deployment is desired anyhow. TIMM’s model format is not a standard PyTorch model and therefore can’t be loaded with e.g. libtorch easily. To load a TIMM model the timm.models.create_model function has to be used. The result is a normal PyTorch model that can be exported easily. If ONNX export is desired, then torch.onnx.export does the job.

A little while ago, I released this export script to export TIMM models to ONNX. Just download it, set the folder_path and export all models to ONNX or select specific models and features using the arguments:

python export.py --help
usage: export.py [-h] [--single_model] [--model_name MODEL_NAME] [--folder_path FOLDER_PATH] [--best_model_only] [--last_model_only] [--input_dim INPUT_DIM] [--num_classes NUM_CLASSES]

optional arguments:
  -h, --help            show this help message and exit
  --single_model        Flag to converts a single model (use full path + filename as --folder_path; input dimensions (--input_dim) and number of classes (--num_classes) are required)
  --model_name MODEL_NAME
                        model name (e.g. mobilenetv3_large_100) required if single_model
  --folder_path FOLDER_PATH
                        Path of folder containing training output (incl. yaml files)
  --best_model_only     Flag to export model_best.pth.tar in training folder only
  --last_model_only     Flag to export last.pth.tar in training folder only
  --input_dim INPUT_DIM
                        Input shape of model assuming square input (example: --input_dim 224); if not provided, input dim is derived from folder name
  --num_classes NUM_CLASSES
                        Number of output classes