$ show lucidrains/vit-pytorch
vit-pytorch is a comprehensive PyTorch library implementing dozens of Vision Transformer (ViT) architectures and extensions, along with self-supervised pretraining methods (MAE, DINO, SimMIM, MP3, LeJEPA) and video transformers (ViViT, NaViT-3D, CCT-3D). It is a research-oriented codebase offering a consistent API across models, all accepting standard hyperparameters like image_size, patch_size, depth, and heads. The library covers a wide spectrum of attention mechanisms (cross-attention, shifted windows, channel attention, sparse gating), tokenization strategies (patch merging, region-based, t2t, SPT), and training utilities (recorders, extractors, distillation wrappers). While primarily a collection of model implementations, it also includes action transformers for robotics (VAT, VAAT) and N-dimensional ViTs, demonstrating both breadth and depth of recent transformer research. The code is highly modular, relying on shared building blocks such as Attention, FeedForward, and Transformer classes, with heavy but well-organized use of einops for tensor manipulation. Some variants are experimental and may lack the polish of the core models, but the overall architecture is clean and serves as a valuable reference for implementing and comparing ViT variations. The repository is actively maintained, with recent commits and a large user base, but its sheer size and variety mean that not every model is exhaustively tested—core components are covered by basic forward-pass tests, while many advanced variants rely on the underlying patterns being sound.
The LLM identified 5 files worth highlighting for reuse:
Core implementation of ViT with reusable Attention, FeedForward, and Transformer classes. Demonstrates clean pre-LayerNorm residual design, multiple attention modes (full, linear), and efficient use of einops for patch embedding. This is the canonical building block reused across most variants in the library and can be extracted as a general-purpose transformer encoder.
A concise, dependency-light ViT implementation using global average pooling instead of a CLS token and simple sinusoidal positional embeddings. The entire model is under ~100 lines and showcases a clean, readable design. It serves as a great starting point for custom modifications and is easily adaptable to other data modalities (1D, 3D, etc.).
A functional wrapper that registers forward hooks on attention layers to capture attention maps without modifying the original model. It provides a clean, side-effect-free way to inspect internal activations, which is highly reusable for any transformer-based model for interpretability or debugging.
Similar to Recorder, this wrapper extracts intermediate embeddings or attention outputs via hook registration. It allows users to obtain task-specific representations from any ViT variant without reimplementation, making it a versatile tool for transfer learning and feature engineering.
Adds a decorrelation loss to the ViT training objective, encouraging diverse feature representations across attention heads. This is a self-contained training module that can be applied to any transformer encoder to improve feature disentanglement, demonstrating a non-obvious regularization pattern.