Vision Transformers

Attention learns to see.

⏱ 2 min read▶ Video at 1:36:28🎓 Free

Watch this lesson in the Deep Learning courseWatch on YouTube ↗Full Deep Learning course

For a decade, convolutions ruled vision. Then, in twenty twenty, a Google team asked a bold question. What if we treated an image like a sentence? Chop it into small patches, treat each patch as a word, and feed them into a transformer. It worked. And it showed that attention was far more general than anyone had guessed.

Convolutions ruled for a decade
Treat an image like a sentence
Patches are words

Image patches as tokens

Here's the recipe. Slice an image into a grid of patches. Turn each patch into a vector, just like a word embedding. Add a position, so that the model knows where each patch came from. Add one special token, which will collect the final answer. Then pass everything through normal transformer blocks. Now each patch can attend to any other, even from the opposite corner of the image, from the very first layer.

Patchesa grid of tiles
Embedeach tile to a vector
Positionwhere it came from
Transformerpatches talk
Class tokengives the answer

A vision transformer

This is a small but real one. A convolution, with a stride the size of the patch, is a neat way to cut and embed patches at once. We add a class token, and positions. PyTorch provides the transformer layers. And the head reads out the class token. Run it. It learns well, in a few epochs. Note that vision transformers need a lot of data to shine, since they don't have the built in assumptions of convolutions. That's why they're often pretrained.

vit.pypython

import torch
from torch import nn
from torchvision import datasets, transforms

class ViT(nn.Module):
    def __init__(self, patch=4, dim=64, heads=4, layers=3, classes=10):
        super().__init__()
        self.patch = nn.Conv2d(1, dim, patch, stride=patch)
        n = (28 // patch) ** 2
        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos = nn.Parameter(torch.randn(1, n + 1, dim) * 0.02)
        layer = nn.TransformerEncoderLayer(dim, heads, dim * 2, batch_first=True,
                                           norm_first=True, activation="gelu")
        self.enc = nn.TransformerEncoder(layer, layers)
        self.head = nn.Linear(dim, classes)

    def forward(self, x):
        x = self.patch(x).flatten(2).transpose(1, 2)
        x = torch.cat([self.cls.expand(len(x), -1, -1), x], 1) + self.pos
        return self.head(self.enc(x)[:, 0])

tf = transforms.ToTensor()
train = datasets.FashionMNIST("data", train=True, transform=tf)
test = datasets.FashionMNIST("data", train=False)
X = test.data[:3000].float().unsqueeze(1) / 255
y = test.targets[:3000]

torch.manual_seed(0)
model = ViT()
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
sub = torch.utils.data.Subset(train, range(20000))
dl = torch.utils.data.DataLoader(sub, batch_size=128, shuffle=True)
for epoch in range(3):
    model.train()
    for a, b in dl:
        loss = nn.functional.cross_entropy(model(a), b)
        opt.zero_grad(); loss.backward(); opt.step()
model.eval()
with torch.no_grad():
    print("Accuracy:", (model(X).argmax(1) == y).float().mean().item())

Output

Accuracy: 0.7956666946411133

Summary

Images become sequences of patches, and attention can see. Today's best vision systems mix both ideas. But all of these models need labels, or do they? Next, one of the most important trends in A I, learning without any labels at all.

Key takeaways

Patches act like words
Class token gives the answer
Needs lots of data or pretraining

Next up: Learning Without Labels — self-supervised learning

❮ Serving a Network Learning Without Labels ❯