Models

This page contains the API reference documentation for creating and discovering timm models, their pretrained configs and the checkpoint / Hub helpers, and an overview of the interface shared by timm models. The individual architectures are covered in the model summaries and the per-model pages.

Usage

Models are created by name with create_model. The name selects an architecture registered in timm, optionally w/ a pretrained tag after a . (resnet50.a1_in1k), or a model on the Hugging Face Hub or in a local folder via a source prefix. Keyword arguments that create_model doesn’t use itself pass through the model’s entrypoint function to the model class constructor, so any model specific argument can be set.

import torch
import timm

# Random init (pretrained=False is the default). The tag selects the pretrained cfg, and the weights loaded with
# pretrained=True. Without a tag the architecture's default tag is used.
model = timm.create_model('vit_tiny_patch16_224.augreg_in21k_ft_in1k')
print(model.pretrained_cfg['mean'], model.pretrained_cfg['crop_pct'])  # (0.5, 0.5, 0.5) 0.9

# Model arguments pass through to the model class, a 10 class head, single channel input and regularization.
model = timm.create_model('resnet50', num_classes=10, in_chans=1, drop_rate=0.2, drop_path_rate=0.05)

# Discover models, filters are fnmatch wildcards.
timm.list_models('convnext*')  # architecture names
timm.list_pretrained('resnet50.*')  # architecture.tag names w/ pretrained weights
timm.list_models(module='vision_transformer', pretrained=True)

# Image preprocessing to match the model's pretrained cfg (input size, interpolation, crop, mean / std).
model = timm.create_model('vit_tiny_patch16_224.augreg_in21k_ft_in1k')
data_config = timm.data.resolve_model_data_config(model)
transform = timm.data.create_transform(**data_config, is_training=False)

x = torch.randn(2, 3, 224, 224)
features = model.forward_features(x)  # unpooled features, (2, 197, 192)
pre_logits = model.forward_head(features, pre_logits=True)  # pooled features, (2, 192)
logits = model.forward_head(features)  # same as model(x), (2, 1000)

model.reset_classifier(5)  # replace the classifier w/ a new 5 class one
model.set_input_size(img_size=(384, 384))  # resample the position embedding for a new resolution
logits = model(torch.randn(2, 3, 384, 384))  # (2, 5)

The training script (train.py) exposes the common arguments as --model, --pretrained, --num-classes, --in-chans, --gp, --drop, --drop-path, --pretrained-path, --initial-checkpoint and --model-kwargs key=value ... for any other model argument. See the quickstart for an introduction.

Model names

The Hub and local sources read the config even w/ pretrained=False, weights are only loaded w/ pretrained=True. A name w/o a source prefix that contains /, \ or : raises a ValueError instead of falling back to a registry model of the same name. hf_hub: is a deprecated alias of hf-hub:.

Common arguments

• pretrained loads the weights of the pretrained cfg, from state_dict or file if set, otherwise from the Hub (hf_hub_id) when huggingface_hub is installed, or the url. If num_classes differs from the pretrained cfg, the classifier weights are discarded (left at random init). If in_chans differs from the source checkpoint’s channel count (input_size[0] in the pretrained cfg), the first conv weights are adapted where supported, so a pretrained model can be created for fine-tuning in one call.
• num_classes sets the classifier size, 0 removes the classifier. global_pool sets the pooling type, the options are model specific (e.g. avg, max, avgmax, catavgmax for most convnets, token, avg, map for ViT), '' disables pooling. in_chans sets the input channels. The num_classes and in_chans defaults come from the pretrained cfg.
• drop_rate (classifier dropout) and drop_path_rate (stochastic depth) are supported by most models.
• img_size is an argument of ViT style and other models w/ a fixed input size, they’re built for the pretrained cfg input_size unless it’s set. The ViT and EVA families also accept dynamic_img_size=True to resample the position embedding to the input size on every forward pass.
• features_only (w/ out_indices, output_stride, feature_cls) wraps the model as a feature backbone returning feature maps, see feature extraction.
• pretrained_cfg replaces the pretrained cfg (a tag, dict or PretrainedCfg), pretrained_cfg_overlay overrides some of its fields, e.g. pretrained_cfg_overlay=dict(file='weights.pth') loads pretrained=True weights from a local file, w/ the classifier / input conv adaptation above (as --pretrained-path does in train.py).
• checkpoint_path loads a checkpoint w/ load_checkpoint after the model is created. It’s a strict load w/o adaptation, EMA weights are used if the checkpoint has them.
• cache_dir overrides the download cache for pretrained weights and Hub configs.
• device and dtype are passed to the model class, params are created on the device and in the dtype directly. device='meta' creates a model w/o allocating memory for its weights.
• scriptable, exportable and no_jit set layer config flags for TorchScript scripting, export / tracing, and to disable jit scripted layers.

Arguments set to None are dropped before reaching the model, so the model defaults apply. Arguments the model class doesn’t accept raise a TypeError in most models, a few ignore unknown arguments.

Pretrained config

Architectures have one or more pretrained cfgs, one per tag, registered w/ the model. A cfg describes the weights (location, license, origin) and how the model was trained and should be evaluated w/ them. Models created by create_model have their cfg as a dict in model.pretrained_cfg (model.default_cfg is an alias), the registry versions are returned as PretrainedCfg instances by get_pretrained_cfg.

The input fields are what resolve_data_config / resolve_model_data_config resolve to a data config for the dataset transforms and loader. The cfg describes the source checkpoint and isn’t updated when the model is created or modified differently (num_classes, in_chans, img_size, set_input_size()). The model’s current input channels, size and fixed size constraint are tracked separately: get_model_input_config returns them, and resolve_input_data_config resolves a data config against them (as the training, validation and inference scripts do). get_model_args returns the arguments to rebuild the model as it is now, which save_for_hf writes to the exported config. See adapt and export models for details.

Model interface

timm models share a common interface beyond forward(): the forward_features() / forward_head() split, classifier reset, intermediate features, input resizing and training hooks such as gradient checkpointing and parameter grouping. It is described in the timm Model API, including which model families support which parts.

Checkpoints and the Hub

load_checkpoint loads a checkpoint file (PyTorch or .safetensors) into a model, load_state_dict returns its state dict. Both pick the EMA weights of a train.py checkpoint by default (use_ema) and remove module. (DDP) and _orig_mod. (torch.compile) key prefixes. resume_checkpoint restores the model, optimizer and loss scaler state of a train.py checkpoint and returns the epoch to resume from.

save_for_hf writes the weights and a config.json (architecture, pretrained cfg, model_args, label_names) to a folder, push_to_hf_hub uploads them to a Hub repo w/ a generated model card if the repo has none. The model is then created w/ hf-hub: or local-dir:, see sharing and loading models.

import timm
from timm.models import save_for_hf

model = timm.create_model('resnet18', num_classes=4)
# ... fine-tune
save_for_hf(model, 'my_resnet18', model_config=dict(label_names=['a', 'b', 'c', 'd']), safe_serialization=True)
model = timm.create_model('local-dir:my_resnet18', pretrained=True)  # 4 classes, weights and labels restored
print(model.pretrained_cfg['label_names'])

Custom models

A model entrypoint function decorated w/ register_model is added to the registry, so it can be created w/ create_model and found w/ list_models. Entrypoints in timm build their model w/ build_model_with_cfg, which handles the pretrained cfg, pretrained weight loading and features_only.

API

Model creation

>>> from timm import create_model

>>> # Create a MobileNetV3-Large model with no pretrained weights.
>>> model = create_model('mobilenetv3_large_100')

>>> # Create a MobileNetV3-Large model with pretrained weights.
>>> model = create_model('mobilenetv3_large_100', pretrained=True)
>>> model.num_classes
1000

>>> # Create a MobileNetV3-Large model with pretrained weights and a new head with 10 classes.
>>> model = create_model('mobilenetv3_large_100', pretrained=True, num_classes=10)
>>> model.num_classes
10

>>> # Create a Dinov2 small model with pretrained weights and save weights in a custom directory.
>>> model = create_model('vit_small_patch14_dinov2.lvd142m', pretrained=True, cache_dir="/data/my-models")
>>> # Data will be stored at */data/my-models/models--timm--vit_small_patch14_dinov2.lvd142m/*

Model discovery

Pretrained config

Current model input and export

Checkpoints and Hub

Custom models