Vision Transformers

For ten years, convolutions ruled vision. Then, in twenty twenty, researchers tried something bold. They cut an image into patches, treated each patch like a word, and fed them to a transformer. It worked, and it changed the field.

• Ten years of convolutions
• Treat patches like words
• It changed the field

An image as a sentence

Let's cut up a photo. We take a two hundred and twenty four pixel image, and slice it into sixteen pixel patches. Run it. That gives one hundred and ninety six patches. Each is flattened into a vector, then attention lets any patch look at any other, from the very first layer. A convolution can only see nearby pixels. This sees the whole picture straight away.

import torch
from skimage import data

image = torch.tensor(data.astronaut()).permute(2, 0, 1)[None].float()
image = torch.nn.functional.interpolate(image, 224)

size = 16
patches = image.unfold(2, size, size).unfold(3, size, size)
patches = patches.reshape(1, 3, -1, size, size).permute(0, 2, 1, 3, 4)

print("Image:", list(image.shape))
print("Patches:", list(patches.shape))
print("Each patch has", 3 * size * size, "numbers")

Image: [1, 3, 224, 224]
Patches: [1, 196, 3, 16, 16]
Each patch has 768 numbers

A pretrained ViT

Here's a real vision transformer, trained on a million images. It works just like any other model. Run it. It has around eighty six million weights, and it recognises the cat correctly. Transformers need a lot of data to shine, so we nearly always start from pretrained ones. Modern systems also mix the two, with convolutions for local detail, and attention for the bigger picture.

import torch
from PIL import Image
from torchvision import models
from skimage import data

w = models.ViT_B_16_Weights.DEFAULT
model = models.vit_b_16(weights=w).eval()
x = w.transforms()(Image.fromarray(data.chelsea())).unsqueeze(0)

with torch.no_grad():
    p = model(x).softmax(1)[0].topk(2)

print("Weights:", round(sum(q.numel() for q in model.parameters()) / 1e6), "million")
for s, i in zip(p.values, p.indices):
    print(f"{w.meta['categories'][i]:<16} {s:.2f}")

Weights: 87 million
Egyptian cat     0.66
tiger cat        0.10

Summary