vit-trainer
A small, educational Python package for fine-tuning Vision Transformer models with PyTorch, reaching 97.65% accuracy on CIFAR-10 with a simpler API than general model-zoo libraries.
vit-trainer focuses specifically on fine-tuning ViT models rather than covering the full range of architectures that libraries such as timm or transformers support, trading breadth for a simple, readable, production-usable API.
It supports mixed-precision training with the AdamW optimizer and cosine annealing with warmup, attention visualization for interpretability, and ONNX export for deployment, usable either from a command-line interface (vit-train) or a Python API for training, evaluation, prediction and metric computation.
It installs with pip, with optional extras for a Gradio demo or ONNX export, and can also be installed from source.