# vit-pytorch

> Readable PyTorch implementations of the Vision Transformer and its variants

- Category: Multimodal Generation / Vision & Understanding
- Language: Python
- Licence: MIT
- Repository: https://github.com/lucidrains/vit-pytorch
- Canonical page: https://ai-tldr.dev/tools/vit-pytorch/

## In short

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. It sits in the Vision & Understanding category of the Multimodal Generation tools and is open-source, written in Python and licensed MIT.

## 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.

1. Install

   One pip package covers every variant in the repository.

```bash
pip install vit-pytorch
```

2. Build and run a Vision Transformer

   image_size must be divisible by patch_size, and the resulting patch count must be greater than 16. The model returns one logit vector per image.

```python
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)
```

3. Swap in a variant or a self-supervised recipe

   Every architecture in the README is importable from its own module under vit_pytorch, and the self-supervised wrappers take a constructed ViT as their argument. If you need pretrained weights rather than a from-scratch model, the README points to Ross Wightman's pytorch-image-models.

## 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