autotimm 0.7.35

Train state-of-the-art vision models with minimal code.

Production-ready computer vision framework powered by timm and PyTorch Lightning.

Overview · Features · Quick Start · Tasks · Architecture · Docs · Contributing · License

Overview

AutoTimm is a production-ready computer vision framework that combines timm (1000+ pretrained models) with PyTorch Lightning. Train image classifiers, object detectors, and segmentation models with any timm backbone using a simple, intuitive API.

Whether you're training an image classifier, an object detector, or a segmentation model, AutoTimm handles:

• Automated model training with auto-tuning of learning rate and batch size
• 4 vision tasks — classification (single & multi-label), object detection, semantic & instance segmentation
• 1000+ backbones — ResNet, EfficientNet, ViT, ConvNeXt, Swin, and more via timm + HuggingFace Hub
• Model interpretation — GradCAM, Integrated Gradients, Attention Rollout, and more
• Production export — TorchScript and ONNX for deployment anywhere

From prototype to production in minutes, not hours.

How It Compares

Features

Quick Start

Installation

pip install autotimm

Everything included: PyTorch, timm, PyTorch Lightning, torchmetrics, albumentations, and more.

pip install autotimm[tensorboard]  # TensorBoard
pip install autotimm[wandb]        # Weights & Biases
pip install autotimm[mlflow]       # MLflow
pip install autotimm[onnx]         # ONNX export (onnx + onnxruntime + onnxscript)
pip install autotimm[all]          # Everything

Your First Model

import autotimm as at  # recommended alias
from autotimm import AutoTrainer, ImageClassifier, ImageDataModule, MetricConfig

# Data
data = ImageDataModule(
    data_dir="./data",
    dataset_name="CIFAR10",
    image_size=224,
    batch_size=64,
)

# Metrics
metrics = [
    MetricConfig(
        name="accuracy",
        backend="torchmetrics",
        metric_class="Accuracy",
        params={"task": "multiclass"},
        stages=["train", "val", "test"],
        prog_bar=True,
    )
]

# Model — try efficientnet_b0, vit_base_patch16_224, hf-hub:timm/resnet50.a1_in1k, etc.
model = ImageClassifier(backbone="resnet18", num_classes=10, metrics=metrics, lr=1e-3)

# Train with auto-tuning (finds optimal LR and batch size automatically)
trainer = AutoTrainer(max_epochs=10)
trainer.fit(model, datamodule=data)

Auto-tuning is enabled by default. Disable with tuner_config=False for manual control.

Command-Line Interface

Train from YAML configs — no Python scripts needed:

autotimm fit --config config.yaml
autotimm test --config config.yaml --ckpt_path best.ckpt

model:
  class_path: autotimm.ImageClassifier
  init_args:
    backbone: resnet18
    num_classes: 10

data:
  class_path: autotimm.ImageDataModule
  init_args:
    dataset_name: CIFAR10
    data_dir: ./data
    batch_size: 32
    image_size: 224

trainer:
  max_epochs: 10
  accelerator: auto

Task Examples

Image Classification

from autotimm import ImageClassifier

model = ImageClassifier(
    backbone="efficientnet_b0",  # or "hf-hub:timm/resnet50.a1_in1k"
    num_classes=10,
    metrics=metrics,
)

trainer = AutoTrainer(max_epochs=10)
trainer.fit(model, datamodule=data)

Object Detection with YOLOX

from autotimm import YOLOXDetector, DetectionDataModule

model = YOLOXDetector(
    model_name="yolox-s",  # nano, tiny, s, m, l, x
    num_classes=80,
    lr=0.01,
    optimizer="sgd",
    scheduler="yolox",
    total_epochs=300,
)

trainer = AutoTrainer(max_epochs=300, precision="16-mixed")
trainer.fit(model, datamodule=DetectionDataModule(data_dir="./coco", image_size=640))

Semantic Segmentation

from autotimm import SemanticSegmentor, SegmentationDataModule

model = SemanticSegmentor(
    backbone="resnet50",
    num_classes=19,
    head_type="deeplabv3plus",
    loss_type="combined",  # CE + Dice for better boundaries
)

data = SegmentationDataModule(data_dir="./cityscapes", format="cityscapes", image_size=512)

trainer = AutoTrainer(max_epochs=100)
trainer.fit(model, datamodule=data)

Instance Segmentation

from autotimm import InstanceSegmentor, InstanceSegmentationDataModule

model = InstanceSegmentor(backbone="resnet50", num_classes=80, mask_loss_weight=1.0)

trainer = AutoTrainer(max_epochs=100)
trainer.fit(model, datamodule=InstanceSegmentationDataModule(data_dir="./coco"))

Multi-Label Classification

from autotimm import ImageClassifier, MultiLabelImageDataModule, MetricConfig

data = MultiLabelImageDataModule(
    train_csv="train.csv", image_dir="./images", val_csv="val.csv",
    image_size=224, batch_size=32,
)
data.setup("fit")

model = ImageClassifier(
    backbone="resnet50",
    num_classes=data.num_labels,
    multi_label=True,
    threshold=0.5,
    metrics=[MetricConfig(
        name="accuracy", backend="torchmetrics", metric_class="MultilabelAccuracy",
        params={"num_labels": data.num_labels}, stages=["train", "val"], prog_bar=True,
    )],
)

trainer = AutoTrainer(max_epochs=10)
trainer.fit(model, datamodule=data)

Model Interpretation

from autotimm.interpretation import (
    GradCAM, GradCAMPlusPlus, IntegratedGradients,
    SmoothGrad, AttentionRollout, AttentionFlow,
    ExplanationMetrics, InteractiveVisualizer,
)

# Explain a prediction
explainer = GradCAM(model)
heatmap = explainer.explain(image, target_class=5)
explainer.visualize(image, heatmap, save_path="gradcam.png")

# Quantitative evaluation
metrics = ExplanationMetrics(model, explainer)
deletion = metrics.deletion(image, target_class=5, steps=50)
insertion = metrics.insertion(image, target_class=5, steps=50)

# Interactive comparison
viz = InteractiveVisualizer(model)
viz.compare_methods(image, {
    'GradCAM': GradCAM(model),
    'IntegratedGradients': IntegratedGradients(model),
}, save_path="comparison.html")

Model Export

from autotimm import ImageClassifier, export_to_torchscript, export_to_onnx
import torch

model = ImageClassifier.load_from_checkpoint("model.ckpt")
example_input = torch.randn(1, 3, 224, 224)

# TorchScript
export_to_torchscript(model, "model.pt", example_input=example_input)

# ONNX
export_to_onnx(model, "model.onnx", example_input=example_input)

Loading from Checkpoints

All task classes inherit PyTorch Lightning's load_from_checkpoint. Most parameters are auto-restored from saved hyperparameters — only ignored parameters (non-serializable objects) need to be re-supplied.

import autotimm as at

# Basic — works for all tasks (ignored params default to None)
model = at.ImageClassifier.load_from_checkpoint("checkpoint.ckpt")

# For inference on CPU
model = at.ImageClassifier.load_from_checkpoint("checkpoint.ckpt", map_location="cpu")
model.eval()

# Override any saved hyperparameter
model = at.ImageClassifier.load_from_checkpoint("checkpoint.ckpt", num_classes=20)

# Disable torch.compile for faster loading (recommended for inference/export)
model = at.ImageClassifier.load_from_checkpoint("checkpoint.ckpt", compile_model=False)

# Resume training with custom metrics and loss
model = at.ObjectDetector.load_from_checkpoint(
    "checkpoint.ckpt",
    metrics=[at.MetricConfig(name="map", backend="torchmetrics", metric_class="MeanAveragePrecision")],
    cls_loss_fn="focal",
    reg_loss_fn="giou",
)

Note: backbone, num_classes, and all other constructor arguments are saved as hyperparameters and restored automatically — no need to re-specify them.

Architecture

graph TD
    subgraph API["Public API · autotimm"]
        Init["__init__.py"] --> Tasks
        Init --> Data
        Init --> Interpret
        Init --> Export
    end

    subgraph Core["Core · autotimm.core"]
        Backbone["backbone.py · 1000+ timm models"]
        Metrics["metrics.py · torchmetrics"]
        Loggers["loggers.py · TB / MLflow / W&B / CSV"]
        Utils["utils.py · seeding, optimizers"]
        Logging["logging.py · loguru"]
    end

    subgraph Tasks["Tasks · autotimm.tasks"]
        Cls["ImageClassifier"]
        Det["ObjectDetector"]
        Seg["SemanticSegmentor"]
        Inst["InstanceSegmentor"]
        YOLOX["YOLOXDetector"]
    end

    subgraph Data["Data · autotimm.data"]
        DM["DataModules"] --> DS["Datasets"]
        DM --> TX["Transforms"]
    end

    subgraph Interpret["Interpretation · autotimm.interpretation"]
        Methods["6 Methods"] --> Viz["Visualization"]
        Methods --> QM["Quality Metrics"]
    end

    subgraph Export["Export · autotimm.export"]
        JIT["TorchScript"]
        ONNX["ONNX"]
    end

    Tasks --> Core
    Data --> Core

Documentation & Examples

Testing

# Run all tests
pytest tests/ -v

# Specific modules
pytest tests/test_classification.py
pytest tests/test_yolox.py
pytest tests/test_interpretation.py

# With coverage
pytest tests/ --cov=autotimm --cov-report=html

Contributing

Contributions, issues, and feature requests are welcome! See the issues page to get started.

1. Fork the repository
2. Create a feature branch — git checkout -b feat/your-feature
3. Commit your changes — git commit -m 'feat: add your feature'
4. Push to the branch — git push origin feat/your-feature
5. Open a Pull Request

git clone https://github.com/theja-vanka/AutoTimm.git && cd AutoTimm
pip install -e ".[dev,all]"
pytest tests/ -v

Citation

@software{autotimm,
  author = {Krishnatheja Vanka},
  title = {AutoTimm: Automatic PyTorch Image Models},
  url = {https://github.com/theja-vanka/AutoTimm},
  year = {2026},
}

License

Distributed under the Apache License 2.0. See LICENSE for details.