A Vision Transformer (ViT) adapts the Transformer architecture to images. Instead of reading characters or words, it reads an image as a sequence of smaller patches, allowing attention to model relationships across the image. section 9.9
A ViT classification pipeline works as follows:
- Splits an image into a grid of patches, flattens them, and embeds each patch in a high-dimensional space. section 9.9
- Prepends a learnable [CLS] token and adds learnable positional embeddings so the model retains patch locations. section 9.9
- Passes the resulting [B, N+1, D] sequence through Transformer encoder blocks using self-attention and multiheaded attention. section 9.9
- Uses only the [CLS] output for classification because it aggregates information attended to from the other image patches. section 9.9
Loading a pretrained ViT
Torchvision can instantiate a ViT-B/16 with ImageNet-1K weights: vit = models.vit_b_16(weights=models.ViT_B_16_Weights.IMAGENET1K_V1). The model contains 88.6 million parameters and supports image classification. section 2.1.3
Figure 9.14 Image converted to image patches, which are then flattened and fed into the model
Follow the patch-conversion diagram to see how an image becomes flattened image-token inputs. section 9.9
Figure 9.15 Vision Transformer architecture. The image tokens are passed through a patch embedding layer, and then positional information is added before being fed into the transformer encoder blocks. Only one single output [CLS] token is used in the final output head for classification.
Inspect the architecture diagram showing patch embedding, positional information, encoder blocks, and the final [CLS] classification output. section 9.9
