lucidrains/vit-pytorch: A ViT Implementation You Assemble Yourself
Implementation of Vision Transformer, a simple way to achieve SOTA in vision classification with only a single transformer encoder, in Pytorch
At a glance
- What is it?
- vit-pytorch collects dozens of Vision Transformer variants behind one import, but it ships no pretrained weights. Here is what it does, how to install it, and where it stops.
- Who is it for?
- Adopt vit-pytorch when you need a readable, MIT-licensed Vision Transformer implementation to modify, train from scratch, or use as a reference while reproducing a paper. Do not adopt it if you need pretrained weights or a drop-in classifier, because the README points to Ross Wightman's repository and the official Jax repository for those.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 11 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
What vit-pytorch actually solves
The repository implements the Vision Transformer described in the paper linked from its README, which the README summarizes as a simple way to achieve state of the art in vision classification with only a single transformer encoder. The audience is narrow and specific: engineers who want to read, modify or train a ViT in PyTorch rather than call a hosted model. The README states the intent plainly, saying there is not much to code here but that laying it out helps expedite the attention revolution. That framing matters. This is a reference implementation plus a catalog, not a model zoo. Anyone arriving from a search for a pretrained checkpoint will be redirected by the README itself, which points to Ross Wightman's pytorch-image-models for pretrained models and to the official Jax repository for the original work.
One encoder, many variants behind the same import
The core mechanism is the standard ViT pipeline: split the image into patches of patch_size, project them, add positional information, run them through depth transformer blocks with heads attention heads and an mlp_dim feedforward layer, and produce num_classes logits. The constructor exposes image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels, dropout, emb_dropout and pool, where pool selects either cls token pooling or mean pooling. Two constraints are enforced by the parameters themselves: image_size must be divisible by patch_size, and the resulting patch count n = (image_size // patch_size) ** 2 must be greater than 16. Rectangular inputs are handled by making image_size the maximum of width and height. Around that core, the README documents a long list of variants, including SimpleViT, NaViT, Deep ViT, CaiT, Token-to-Token ViT, CCT, Cross ViT, PiT, LeViT, CvT, Twins SVT, CrossFormer, RegionViT, ScalableViT, SepViT, MaxViT, NesT, MobileViT, XCiT, Masked Autoencoder, Dino and EsViT. They live in the same package and follow the same constructor style, which is the practical value: switching architecture is an import change, not a rewrite. NaViT is the most interesting divergence, packing images of multiple resolutions into a single batch with masking and factorized 2d positional encodings. It requires you to place images in the same batch element so the sequence does not exceed the maximum length for masked self-attention, unless you pass group_images=True with group_max_seq_len, in which case the framework groups them for you. A nested tensor flavor exists for PyTorch 2.5 and later, importing from vit_pytorch.na_vit_nested_tensor.
Installing vit-pytorch and running a first forward pass
Installation is a single pip command, as the README's Install section shows. The package name on PyPI is vit-pytorch.
$ pip install vit-pytorchThe distribution declares dependencies on einops>=0.8.2, torch>=2.5 and torchvision, and requires Python 3.8 or newer. The torch>=2.5 floor is worth noticing: if your environment is pinned to an older PyTorch, pip will try to upgrade it. The README's Usage section gives this example, which builds a model and runs one random tensor through it.
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)The output shape is (1, 1000), one logit vector per image. If you would rather follow the simplification proposed in the update the README cites, SimpleViT drops the CLS token for global average pooling and removes dropout.
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)The repository also carries examples/cats_and_dogs.ipynb and a top-level train_vit_decorr.py, which are the places to look for a training loop rather than the README, which stops at forward passes. The tests directory and the pytest configuration in pyproject.toml point at tests and the repository root with python_files set to test_*.py and *_test.py.
No pretrained weights, and that is the main limitation
The README does not offer a checkpoint, a download URL or a model hub integration. It explicitly redirects readers to Ross Wightman's repository for a PyTorch implementation with pretrained models and to the official Jax repository for the original. If your task is image classification on a modest dataset, this package gives you an untrained architecture and nothing else. You supply the data pipeline, the augmentation, the optimizer schedule and the compute. The README's own description of the SimpleViT update mentions batch sizes of 1024 and RandAugment and MixUp augmentations, which is a useful signal about the regime these models expect and a warning about what a single GPU with a small batch will produce. A second limitation is version churn. The released versions in the pyproject.toml metadata show 1.26.6, while the release notes list 1.17.8 from 2026-02-11, 1.17.7 from 2026-02-04 and 1.17.6 from 2026-01-28. The package is classified as Development Status 4 - Beta, so pinning a version is sensible for anything you intend to reproduce. The last push to the repository was on 2026-09-05, and the repository is not archived, so the code is being touched, but a moving target is still a moving target.
How it differs from pytorch-image-models and Hugging Face
The obvious comparison is Ross Wightman's pytorch-image-models, which the README itself names as the place to go for pretrained models. The difference in approach is stark. timm is a model library with weight loading, a consistent create_model interface and a large set of pretrained checkpoints, aimed at people who want a working classifier today. vit-pytorch is a set of implementations, aimed at people who want to see the code, change the attention, or reproduce a paper variant that timm may not carry. The second comparison is Hugging Face transformers, which centers on pretrained checkpoints and a trainer abstraction. Neither comparison makes vit-pytorch redundant, but neither is it a substitute. If you are choosing between them, the question is whether you need weights or whether you need readable, editable architecture code. vit-pytorch answers only the second.
Licence and what upgrading costs you
The project is MIT licensed, and the pyproject.toml declares the license as a file reference to LICENSE. MIT is permissive: it allows commercial use and modification with attribution and no warranty. That is a genuine advantage over research code released under non-commercial terms, and it is the reason a company can vendor this into an internal training pipeline without a legal review cycle. It is not legal advice, and the LICENSE file is the authority. On upgrade cost, the dependency floor of torch>=2.5 is the practical constraint. A minor release that raises that floor forces a PyTorch upgrade across your environment, and PyTorch upgrades are rarely free for a training stack. The version gap between the metadata and the release notes suggests the package version and the release feed are not perfectly in step, so pinning an exact version and reading the release notes before bumping is the safe path. There is no documented rollback procedure in the README, which means your rollback plan is a pinned requirement and a lockfile.
Editorial conclusion
Adopt vit-pytorch when you need a readable, MIT-licensed Vision Transformer implementation to modify, train from scratch, or use as a reference while reproducing a paper. Do not adopt it if you need pretrained weights or a drop-in classifier, because the README points to Ross Wightman's repository and the official Jax repository for those. Before committing, check that torch>=2.5 is acceptable in your environment, that your image_size is divisible by patch_size, and that the number of patches stays above 16.
Frequently asked questions
What is a ViT model?
A Vision Transformer applies a transformer encoder to image patches instead of tokens from text. The README describes it as a simple way to achieve state of the art in vision classification with only a single transformer encoder.
How does ViT work?
The image is divided into patches of patch_size, and the number of patches is (image_size // patch_size) ** 2, which must be greater than 16. Those patches pass through depth transformer blocks with heads attention heads and an mlp_dim feedforward layer, and the model outputs num_classes logits.
What is a ViT?
In this repository it is the model implemented by the ViT class in vit_pytorch, with configurable image_size, patch_size, num_classes, dim, depth, heads and mlp_dim. The README links the original paper and the official Jax implementation.
Community notes