Guide · Open source

How to use vit-pytorch: build a Vision Transformer in PyTorch

Updated 2026-10-29 · written with AI from the project's own documentation and checked against our catalogue data

vit-pytorch is an MIT-licensed Python library that implements the Vision Transformer, an image classifier built from a single transformer encoder, along with many variants from research papers. This guide shows you how to install it, create a model, and use the main variants. It assumes you already know Python and PyTorch, since the usage examples depend on both.

vit-pytorchlucidrains/vit-pytorch · ★ 25.6k · MIT Install options & downloads

Quick start

1Install the package. Run pip install vit-pytorch in an environment that already has Python and PyTorch.
2Import the model. Import ViT or SimpleViT from vit_pytorch.
3Set the model size. Choose image_size, patch_size, num_classes, dim, depth, heads and mlp_dim, making sure image_size is divisible by patch_size.
4Check the patch count. Confirm that (image_size // patch_size) ** 2 is greater than 16.
5Run a test batch. Create a random tensor shaped (batch, 3, height, width) and pass it to the model to confirm the output shape.
6Train your model. Add your own dataset and training loop in PyTorch, then use the variant that fits your needs.

Install vit-pytorch with pip

The library is published on PyPI, so you can install it with pip. You need Python and PyTorch available in your environment before you run the examples.

Install Python and PyTorch first, as the library depends on them.
Run the install command in the same environment where you will write your code.

$ pip install vit-pytorch

Create a basic ViT image classifier

Import ViT from vit_pytorch and set the core settings. The image_size must be divisible by patch_size, and the number of patches must be greater than 16. Pass a batch of images shaped (batch, channels, height, width) to the model to get predictions shaped (batch, num_classes).

image_size: the maximum of width and height for rectangular images.
patch_size: the size of each patch; image_size must be divisible by it.
num_classes, dim, depth, heads and mlp_dim control the number of classes, embedding size, transformer blocks, attention heads and feed-forward width.
dropout and emb_dropout set dropout rates between 0 and 1.

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)

Understand the main parameters

Most settings are straightforward, but two rules matter most. The patch count is (image_size // patch_size) ** 2, and that value must be greater than 16. The channels setting defaults to 3, and pool can be cls token pooling or mean pooling.

channels: defaults to 3, the number of image channels.
dropout: defaults to 0.
emb_dropout: defaults to 0.
pool: either cls or mean pooling.

Use SimpleViT for a simplified variant

SimpleViT comes from an update by some of the original authors. It uses 2D sinusoidal positional embedding and global average pooling with no CLS token, and it drops dropout. The README notes that a simple linear layer at the end is not significantly worse than the original MLP head.

You import it the same way as ViT and pass the same kind of arguments, without dropout settings in the example.

import torch
from vit_pytorch import SimpleViT

v = SimpleViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 16,
    mlp_dim = 2048
)

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

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

Train with variable-size images using NaViT

NaViT packs images of different resolutions into a single batch. The README says you can pass a list of lists of tensors and place images carefully so the maximum sequence length for self-attention is not exceeded. You can also set group_images = True with group_max_seq_len to let the framework group images automatically.

A nested-tensor flavor is available in vit_pytorch.na_vit_nested_tensor. The README says it requires version 2.5 and skips much of the masking and padding. The README does not state which version of the library this refers to, so check the repository for that detail.

Import from vit_pytorch.na_vit for the standard version.
Import from vit_pytorch.na_vit_nested_tensor for the nested-tensor version.

images = [
    torch.randn(3, 256, 256),
    torch.randn(3, 128, 128),
    torch.randn(3, 128, 256),
    torch.randn(3, 256, 128),
    torch.randn(3, 64, 256)
]

preds = v(
    images,
    group_images = True,
    group_max_seq_len = 64
) # (5, 1000)

Distill a CNN teacher into a ViT

The distillation module lets a convolutional network teach a smaller vision transformer. You wrap a DistillableViT student with DistillWrapper and a teacher model, then call the wrapper with images and labels to get a loss. When training is done, call .to_vit() to get a regular ViT instance back.

The README's example loads a torchvision resnet50 as the teacher.

temperature: the distillation temperature.
alpha: the trade-off between main loss and distillation loss.
hard: whether to use soft or hard distillation.

from vit_pytorch.distill import DistillableViT, DistillWrapper

distiller = DistillWrapper(
    student = v,
    teacher = teacher,
    temperature = 3,
    alpha = 0.5,
    hard = False
)

loss = distiller(img, labels)
loss.backward()

v = v.to_vit()

What people use it for

Building an image classifier in PyTorch with a transformer encoder.
Studying how Vision Transformers work by reading readable code.
Comparing ViT variants such as SimpleViT, NaViT and MaxViT from different papers.
Distilling knowledge from a convolutional teacher such as ResNet50 into a smaller ViT.

Frequently asked questions

How do I install vit-pytorch?

Run pip install vit-pytorch. The library is on PyPI, and the README lists Python and PyTorch as prerequisites.

Does vit-pytorch include pretrained models?

No. The README points to Ross Wightman's pytorch-image-models repository for PyTorch implementations with pretrained models.

Is vit-pytorch free to use?

The repository is licensed under MIT. Check the license file in the repository for the full terms.

What do I need before using vit-pytorch?

You need Python and PyTorch. The usage examples assume familiarity with both, so beginners may find it hard without that background.

What is the difference between ViT and SimpleViT?

SimpleViT is a simplified version from an update by some of the original authors. It uses 2D sinusoidal positional embedding and global average pooling with no CLS token, and it drops dropout.

What are alternatives to vit-pytorch?

The README-listed comparable projects are pytorch-image-models, vit-tensorflow and vision_transformer. The README links pytorch-image-models for pretrained weights and vit-tensorflow as a TensorFlow2 translation.

New agent skills, repos and AI tools — every week

The best new skills and open-source projects, the fastest-growing repos, and one free prompt. Five minutes to read. Unsubscribe anytime.

Our verdict

vit-pytorch suits developers, students and researchers who want readable, modular Vision Transformer code in PyTorch and want to try variants or distillation. If you need pretrained weights or a finished application, pick pytorch-image-models or another ready-to-use project instead.

Install options & downloads