vit-base-patch16-224-in21k

Technical Overview

The google/vit-base-patch16-224-in21k model is a Vision Transformer (ViT‑Base) that treats an image as a sequence of fixed‑size patches and processes them with a transformer encoder, much like BERT processes word tokens. Trained on the massive ImageNet‑21k dataset (≈14 M images, 21 843 classes) at a resolution of 224×224, it learns a rich, generic visual representation that can be reused for downstream tasks such as image classification, retrieval, and feature‑based clustering.

Key capabilities include:

• Patch size 16×16 → 196 tokens per image plus a [CLS] token.
• Pre‑trained pooler for extracting a single‑vector image embedding.
• Fully compatible with 🤗 Transformers in PyTorch, TensorFlow and JAX/Flax.
• Supports image-feature-extraction pipelines out of the box.

Architecture highlights:

• Transformer depth: 12 encoder layers.
• Hidden size: 768 units.
• Feed‑forward dimension: 3072.
• Attention heads: 12, each with a head size of 64.
• Absolute positional embeddings added to the patch embeddings.

The model is deliberately released without any task‑specific classification head – the head was zeroed by the original Google team. This design encourages users to fine‑tune a lightweight classifier on top of the [CLS] token or to use the pooled output as a high‑quality image descriptor for retrieval, clustering, or as input to downstream vision pipelines.

Benchmark Performance

ViT‑Base (224) is evaluated primarily on image‑classification benchmarks. The original paper reports top‑1 accuracy on ImageNet‑1k after fine‑tuning at higher resolutions (384×384), but the raw pre‑trained checkpoint achieves competitive results even at the native 224×224 resolution. Representative numbers from the paper’s tables (2 & 5) are:

• ImageNet‑1k (fine‑tuned, 384×384): ~78 % top‑1 accuracy.
• ImageNet‑21k (pre‑training): ~71 % top‑1 accuracy on the 21k‑class validation set.

These benchmarks matter because they reflect how well the model captures visual semantics at scale. Compared to contemporary CNNs (e.g., ResNet‑50) ViT‑Base offers higher accuracy with a similar parameter count (≈86 M) while providing a more flexible, token‑based representation that can be repurposed for non‑classification tasks. When fine‑tuned on domain‑specific data, ViT‑Base often surpasses ResNet‑101 and matches larger ViT variants (e.g., ViT‑Large) with far less compute.

Hardware Requirements

Running inference with vit-base-patch16-224-in21k is modest by modern transformer standards, but the model’s 86 M parameters still demand a decent GPU for low‑latency applications.

• VRAM for inference: ~4 GB (FP32). Using torch.float16 or torch.bfloat16 reduces this to ~2 GB.
• Recommended GPUs: NVIDIA RTX 3060 or newer, AMD Radeon RX 6700 XT, or any GPU with ≥6 GB VRAM for batch‑size > 1.
• CPU requirements: A recent multi‑core CPU (e.g., Intel i7‑9700K, AMD Ryzen 7 3700X) is sufficient for preprocessing; the heavy lifting remains on the GPU.
• Storage: Model files (weights + config) occupy ~350 MB when stored as .safetensors. Including the processor and tokenizer adds another ~50 MB.
• Performance: On an RTX 3080, a single image passes through the encoder in ~3 ms (FP16). Batch inference scales linearly up to the VRAM limit.

Use Cases

Because the model outputs a high‑dimensional image embedding, it shines in scenarios where a generic visual feature is needed:

• Image classification pipelines: Attach a lightweight linear head to the [CLS] token and fine‑tune on a domain‑specific dataset (e.g., medical imaging, satellite imagery).
• Content‑based image retrieval: Store the pooled embeddings in a vector database and perform nearest‑neighbor search for similarity matching.
• Zero‑shot or few‑shot learning: Use the embeddings as inputs to a prototypical network or a CLIP‑style multimodal model.
• Feature extraction for downstream vision tasks: Feed the embeddings into object detection, segmentation, or video frame analysis pipelines.

Industries that benefit include e‑commerce (product recommendation), digital asset management, autonomous vehicles (scene understanding), and healthcare (radiology image triage). The model integrates seamlessly with the 🤗 Transformers ViTImageProcessor, enabling rapid prototyping in Python, JAX, or TensorFlow.

Training Details

The model was pre‑trained on the ImageNet‑21k dataset (≈14 M images, 21 843 categories). Training was performed on Google’s TPUv3 pods (8 cores) using the following pipeline:

• Image preprocessing: Images resized to 224×224, then normalized per channel with mean = (0.5, 0.5, 0.5) and std = (0.5, 0.5, 0.5).
• Batch size: 4096 images per step (distributed across the 8 TPU cores).
• Learning‑rate schedule: Linear warm‑up for the first 10 k steps, followed by cosine decay.
• Gradient clipping: Global norm ≤ 1 to stabilize training.
• Optimizer: AdamW with weight decay of 0.1 (as per the original ViT codebase).

Fine‑tuning is straightforward: replace the zeroed classification head with a new linear layer, optionally increase the input resolution to 384×384 for better accuracy, and train on the target dataset for a few epochs. The pre‑trained pooler can also be used directly for feature extraction without any additional training.

Related Papers

The model is grounded in two seminal arXiv pre‑prints:

• An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale (Dosovitskiy et al., 2020) – introduces the Vision Transformer architecture.
• Visual Transformers: Token‑based Image Representation and Processing for Computer Vision (Wu et al., 2020) – expands on training strategies and large‑scale pre‑training.

These works demonstrate that a pure transformer, when fed with patch embeddings, can match or exceed the performance of convolutional networks on large‑scale image classification tasks. The vit-base-patch16-224-in21k checkpoint is a direct implementation of the architecture described in the first paper, with training recipes refined in the second.

Licensing Information

The repository tag lists license:apache-2.0, yet the model card’s license field is marked “unknown”. In practice, the underlying weights were originally released under the Apache‑2.0 license by Google and subsequently converted by Ross Wightman, who also respects the same license. Apache‑2.0 is a permissive open‑source license that:

• Allows commercial, academic, and personal use without fee.
• Permits modification, redistribution, and inclusion in proprietary products.
• Requires preservation of the original copyright notice and a copy of the license.
• Mandates a clear attribution to the original authors (Google Research) and the converter (Ross Wightman).

If you distribute a derivative work, you must include a notice such as “Based on Google’s Vision Transformer (ViT‑Base) model, licensed under Apache‑2.0”. No “copyleft” obligations exist, making the model safe for commercial deployment provided the attribution clause is respected.