timm

Loads pre-trained PyTorch vision models from a unified registry (900+ architectures) with automatic weight downloading and caching. Uses a factory pattern with model name resolution to instantiate architectures like ResNet, Vision Transformer, EfficientNet, and proprietary variants. Handles checkpoint loading, device placement, and inference-mode setup in a single call, abstracting away boilerplate PyTorch initialization.

Provides composable image transforms (resize, normalization, augmentation) optimized for vision models with automatic resolution inference from model metadata. Uses PyTorch's torchvision.transforms as a base but adds model-specific defaults (e.g., ImageNet normalization stats, optimal input sizes) and integrates with timm's model registry to auto-configure preprocessing for any loaded model. Supports both training (with augmentation) and inference modes.

Provides a plugin system for registering custom model architectures into the timm registry, enabling them to be loaded via the standard `timm.create_model()` API alongside built-in models. Uses a decorator-based registration pattern that integrates custom models with timm's preprocessing, export, and benchmarking utilities. Supports model composition (combining modules from different architectures) and automatic documentation generation.

Provides a searchable registry of 900+ vision model architectures with filtering by family (ResNet, ViT, EfficientNet), input resolution, parameter count, and training dataset. Exposes model metadata (FLOPs, throughput, accuracy benchmarks) via a programmatic API and CLI. Uses a hierarchical naming convention (e.g., 'resnet50.tv_in1k') to encode architecture, variant, and training source, enabling semantic model selection without manual documentation lookup.

Provides utilities for efficient transfer learning including layer freezing, selective unfreezing, learning rate scheduling per layer group, and checkpoint management. Integrates with PyTorch's optimizer API to enable differential learning rates (e.g., lower LR for early layers, higher for head). Supports both full fine-tuning and adapter-style approaches via selective parameter freezing. Includes utilities for loading partial checkpoints (e.g., pre-trained backbone only) and handling shape mismatches when adapting to new classification heads.

Exports PyTorch models to ONNX, TorchScript, and other inference formats with automatic shape inference and optimization. Handles model-specific export quirks (e.g., handling attention masks in Vision Transformers) and validates exported models against the original PyTorch version. Includes utilities for quantization-aware training (QAT) and post-training quantization (PTQ) to reduce model size for edge deployment.

Provides utilities for efficient batch inference across multiple images with automatic GPU/CPU device placement, mixed precision (fp16/bf16) support, and memory-efficient inference modes. Handles variable-sized inputs by padding or resizing to a common shape. Includes profiling utilities to measure throughput and latency per batch size, enabling automatic batch size selection for hardware constraints.

Provides utilities for combining predictions from multiple models (different architectures, checkpoints, or augmentations) using voting, averaging, or learned weighting strategies. Supports test-time augmentation (TTA) by averaging predictions across multiple augmented versions of the same input. Handles ensemble-specific optimizations like shared preprocessing and batch-level parallelization across ensemble members.

Provides tools for visualizing model predictions, attention maps (for Vision Transformers), and feature activations. Includes gradient-based visualization (Grad-CAM, saliency maps) and attention rollout for understanding which image regions influence predictions. Integrates with timm's model registry to automatically extract attention layers from Vision Transformers and other attention-based architectures.

Provides utilities for distributed training across multiple GPUs (single machine) and multiple nodes (cluster) using PyTorch's DistributedDataParallel (DDP) and automatic mixed precision (AMP). Handles synchronization of batch normalization statistics across devices, gradient accumulation for effective larger batch sizes, and automatic learning rate scaling based on world size. Includes utilities for distributed checkpoint saving and resuming.

Provides tools for benchmarking model inference speed, memory usage, and FLOPs across different batch sizes, input resolutions, and hardware configurations. Includes profiling utilities to identify bottlenecks (compute-bound vs memory-bound), measure throughput (images/sec), and estimate latency percentiles. Integrates with PyTorch profiler to generate detailed performance traces and supports comparison across model families.