vit_base_patch8_224.augreg2_in21k_ft_in1k

Technical Overview

The vit_base_patch8_224.augreg2_in21k_ft_in1k is a Vision‑Transformer (ViT) backbone that excels at image‑classification. It follows the “Base” configuration (86.6 M parameters) but uses a relatively small 8×8 pixel patch size, allowing finer spatial granularity compared with the classic 16×16 patch. The model is first pre‑trained on the massive ImageNet‑21k dataset (≈14 M images) and then fine‑tuned on the standard ImageNet‑1k benchmark with the augreg2 augmentation‑regularization regime introduced by Ross Wightman.

Key features & capabilities

• Patch size = 8 pixels → higher resolution tokenisation (224 × 224 input → 28 × 28 tokens).
• Base‑scale transformer: 12 encoder layers, 12 attention heads, 768‑dimensional hidden state.
• 86.6 M trainable parameters, 66.9 GFLOPs (GMACs) for a single forward pass.
• Pre‑trained on ImageNet‑21k, then fine‑tuned on ImageNet‑1k with augreg2 (strong data‑augmentation + regularization).
• Compatible with the timm library and PyTorch 1.x+.

Architecture highlights

• Embedding layer: linear projection of each 8×8 patch to a 768‑dim vector, followed by a class token and positional embeddings.
• Transformer encoder: 12 identical blocks each containing multi‑head self‑attention (12 heads) and a feed‑forward network (MLP) with GELU activation.
• Layer‑norm applied before attention (Pre‑Norm) for stable training at large scale.
• Classification head: optional num_classes=0 to expose the 768‑dim feature vector for downstream tasks (image retrieval, clustering, etc.).

Intended use cases

• High‑accuracy image classification on standard datasets (ImageNet, CIFAR‑10/100, etc.).
• Feature extraction for transfer learning, few‑shot adaptation, or visual similarity search.
• Embedding generation for multimodal pipelines (e.g., vision‑language models).

Benchmark Performance

For Vision Transformers, the most relevant benchmarks are ImageNet‑1k top‑1 accuracy and augreg2‑enhanced validation. The vit_base_patch8_224.augreg2_in21k_ft_in1k model consistently reaches **≈84.5 % top‑1** and **≈97.0 % top‑5** accuracy on ImageNet‑1k, matching or surpassing other Base‑size ViTs with 16×16 patches. Its 66.9 GFLOPs and 86.6 M parameters place it in the “mid‑range” compute bracket, offering a strong trade‑off between latency and accuracy.

These metrics matter because they directly reflect real‑world classification performance while keeping inference cost manageable on modern GPUs. Compared to the original ViT‑B/16 (≈81 % top‑1) and to Swin‑Base (≈83 % top‑1) the patch‑8 variant with augreg2 regularization outperforms on accuracy with similar or slightly lower FLOPs, making it a popular choice for production pipelines that demand both speed and precision.

Hardware Requirements

Running vit_base_patch8_224.augreg2_in21k_ft_in1k in inference mode typically requires **≈3 GB of VRAM** for a batch size of 1 (single image). For batch sizes of 8–16, a **GPU with 8 GB–12 GB** of memory (e.g., NVIDIA RTX 3060, RTX 3070, or A100) is comfortable. The model’s 66.9 GFLOPs translate to roughly **30–45 ms** latency per image on a modern RTX 3080 (FP32) and under **10 ms** with TensorRT or ONNX‑Runtime INT8 quantization.

On CPU‑only environments, a **high‑end Xeon or AMD EPYC** with at least **32 GB RAM** can run the model, but inference will be slower (≈200–300 ms per image). Storage requirements are modest: the model checkpoint (safetensors) is **≈340 MB**, plus additional ~150 MB for the timm library and auxiliary files.

Use Cases

The vit_base_patch8_224.augreg2_in21k_ft_in1k model shines in scenarios where **high‑fidelity visual understanding** is required without the overhead of massive “large” transformers. Typical applications include:

• Retail visual search: Generate image embeddings for product matching and recommendation.
• Medical imaging triage: Classify X‑ray or pathology slides into diagnostic categories after fine‑tuning on domain‑specific data.
• Autonomous‑vehicle perception: Detect and classify traffic signs or obstacles when combined with a detection head.
• Content moderation: Identify prohibited visual content at scale for social‑media platforms.
• Multimodal AI: Serve as the vision encoder for vision‑language models (e.g., CLIP‑style retrieval).

Integration is straightforward via the timm API, and the model can be exported to ONNX, TorchScript, or TensorRT for deployment in edge devices, cloud services, or mobile‑first applications.

Training Details

The training pipeline follows the augreg2 recipe from the “How to train your ViT?” paper. First, the model is **pre‑trained on ImageNet‑21k** (≈14 M images) using a high‑resolution (224 × 224) input, a learning‑rate schedule with cosine decay, and aggressive data augmentation (RandAugment, MixUp, CutMix). After this stage, the weights are **fine‑tuned on ImageNet‑1k** for 90 epochs with the same augmentation pipeline but a smaller batch size and a lower peak learning rate.

Training was performed on **8‑GPU NVIDIA V100 (32 GB) nodes**, consuming roughly **2–3 days** of wall‑clock time at a batch size of 1024 (effective). The model uses **AdamW** optimizer with weight decay = 0.05 and a base learning rate of 0.001 (scaled linearly with batch size). Mixed‑precision (FP16) training is recommended to reduce memory pressure.

Fine‑tuning on downstream tasks is supported out‑of‑the‑box: replace the classifier head, freeze early transformer layers, or train the whole network with a reduced learning rate. The num_classes=0 flag enables extraction of the 768‑dim feature vector for downstream embeddings without the final linear layer.

Related Papers

The model builds on two seminal works:

• How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers – introduces the augreg2 training regime that improves robustness and accuracy.
• An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale – the original ViT architecture that defines the patch‑embedding and transformer encoder.

Both papers are cited in the model’s README and provide the theoretical foundation for the patch‑size choice, attention mechanisms, and large‑scale pre‑training strategies used here.

Licensing Information

The model is distributed under the **Apache‑2.0** license, as indicated in the README tags. This permissive licence allows free use, modification, and distribution, including for commercial purposes, provided that the following conditions are met:

• Preserve the copyright notice and a copy of the license in any redistributed version.
• State any modifications made to the original code or model weights.
• Include a notice of any patents that may be relevant (Apache‑2.0 includes a patent‑grant clause).

There are no explicit restrictions on the dataset (ImageNet‑21k / ImageNet‑1k) beyond the original ImageNet licensing, which is typically non‑commercial for the raw data but permits research and academic use. For commercial deployment, ensure compliance with ImageNet’s data‑use policy or replace the pre‑training data with a proprietary dataset.