vit-pytorch

// Overview

vit-pytorch is an implementation of the Vision Transformer in PyTorch — the architecture that showed you can reach state-of-the-art image classification with a single transformer encoder over image patches, no convolutions required. The base `ViT` class takes an image size, a patch size and the usual transformer hyper-parameters, and returns class logits for a batch of images.

What makes the package unusual is its breadth. Rather than stopping at the original paper, it collects readable implementations of a long list of follow-up architectures in one importable library: Simple ViT, NaViT, DeepViT, CaiT, T2T-ViT, CCT, CrossViT, PiT, LeViT, CvT, Twins SVT, CrossFormer, RegionViT, ScalableViT, SepViT, MaxViT, NesT, MobileViT, XCiT, 3D ViT and ViViT among others, plus self-supervised recipes including the Masked Autoencoder, masked patch and position prediction, DINO and EsViT, and a distillation wrapper.

The code is written to be read. Each variant is a compact, self-contained module you can import, subclass, or lift into your own model, which is why the repository is so often used as a reference when implementing a paper. It does not ship pretrained weights — for pretrained ViT checkpoints the README points at Ross Wightman's pytorch-image-models, and at the official Jax repository for the original release.

// What it does

• The reference ViT module, configurable by image size, patch size, depth, heads and MLP dimension
• Dozens of architecture variants in one package — Simple ViT, NaViT, CaiT, CvT, LeViT, MaxViT, MobileViT, XCiT, CrossViT and more
• Video and volumetric variants: 3D ViT and ViViT
• Self-supervised training wrappers including Masked Autoencoder, masked patch prediction, DINO and EsViT
• A distillation wrapper for training a ViT from a convolutional teacher
• Helpers for accessing attention maps, useful for visualising what a model attends to

// Getting started

Install the package and instantiate a ViT with your image and patch size. The snippet below is the README's usage example.

pip install vit-pytorch

import torch
from vit_pytorch import ViT

v = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)

img = torch.randn(1, 3, 256, 256)

preds = v(img) # (1, 1000)

Commands and code are distilled from the project's own documentation — always check the official repo for the latest.

// When to use it

• Train an image classifier on your own dataset with a transformer backbone instead of a CNN
• Compare several ViT variants on the same task without reimplementing each paper
• Read a compact reference implementation while implementing a vision-transformer paper yourself
• Pre-train a vision backbone with a self-supervised objective such as the Masked Autoencoder or DINO

// How vit-pytorch compares

vit-pytorch alongside other open-source vision & understanding tools AI/TLDR tracks, ranked by GitHub stars.