Vision Transformer (ViT) — это архитектура, которая буквально произвела революцию в том, как машины «видят» мир.
В этой статье я не просто объясню, что такое ViT — я покажу вам, как создать эту магию своими руками, шаг за шагом, даже если вы никогда раньше не работали с трансформерами для задач с изображениями.
Для начала давайте взглянем на архитектуру Vision Transformer:
Мы напишем код полностью с нуля, а затем обучим модель на датасете CIFAR-10.
Давайте начнём с реализации Patch Embedding:
class PatchEmbedding(nn.Module): def __init__(self, img_size = 32, patch_size = 4, in_channels = 3, embed_dim=256): super().__init__() assert img_size % patch_size == 0, "img_size must be divisible by patch_size" self.patch_size = patch_size self.num_patches = (img_size//patch_size)**2 self.conv = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): x = self.conv(x) #(B, embed_dim, H/patch_size, W/patch_size) x = x.flatten(2).transpose(1, 2) #(B, num_patches, embed_dim) return x
Изображение будет разделено на патчи, и размер каждого патча можно задать с помощью параметра patch_size. При этом изображение не просто разбивается на патчи, но и пропускается через свёрточные ядра (CNN). В итоге мы получаем не просто патчи изображения — а встраивания (эмбеддинги) этих патчей.
Следующий шаг — реализовать самую интересную часть этой модели — механизм внимания (attention).
Q (Query) формально задаёт вопрос от каждого патча к другим патчам, K (Key) показывает, есть ли у каждого патча ответ на этот вопрос, а V (Value) содержит «значения» — фактические данные каждого патча, которые используются для формирования итогового представления.
Предположим, у нас есть X и Y, и мы хотим, чтобы X обращал внимание на Y. В этом случае матрица Query умножается на X, а матрицы Key и Value — на Y. Вместо прямого умножения на матрицы мы используем линейные слои.
attn_probs — это матрицы внимания, которые показывают, насколько токен i должен «обращать внимание» на токен j. Далее мы умножаем их на V, чтобы получить эмбеддинги изображения с учётом весов внимания attn_probs. V фактически хранит значения каждого патча изображения, а attn_probs показывает, сколько информации каждый патч должен получить от остальных патчей.
Вот как работает одна голова внимания; затем значения с всех голов объединяются. Такая конструкция основана на идее, что каждая голова фокусируется на разных аспектах.
class MultiHeadAttention(nn.Module): def __init__(self, dim, num_heads, dropout): super().__init__() assert dim % num_heads == 0, "Embedding dimension must be divisible by number of heads" self.num_heads = num_heads self.head_dim = dim // num_heads self.qkv = nn.Linear(dim, dim * 3, bias = False) self.out = nn.Linear(dim, dim, bias = False) self.scale = 1.0 / (self.head_dim ** 0.5) self.attn_dropout = nn.Dropout(dropout) def forward(self, x, mask = None, return_attn=False): B, num_patches, embed_dim = x.shape qkv = self.qkv(x) # (B, num_patches, 3*embed_dim) qkv = qkv.reshape(B, num_patches, 3, self.num_heads, self.head_dim) qkv = qkv.permute(2, 0, 3, 1, 4) #(3, B, num_heads, num_patches, head_dim) q, k, v = qkv[0], qkv[1], qkv[2] # each (B, num_heads, num_patches, head_dim) #How important it is for token i to pay attention to token j. attn_scores = (q @ k.transpose(-2, -1)) * self.scale #[B, num_heads, N, N] if mask is not None: # mask: (B, 1, N, N) or (1, 1, N, N) attn_scores = attn_scores.masked_fill(mask == 0, float("-inf")) attn_probs = attn_scores.softmax(dim=-1) #[B, num_heads, N, N] attn_probs = self.attn_dropout(attn_probs) attn_output = attn_probs @ v # (B, num_heads, num_patches, head_dim) attn_output = attn_output.transpose(1, 2).reshape(B, num_patches, embed_dim) if return_attn: return self.out(attn_output), attn_probs else: return self.out(attn_output) #(B, num_patches, embed_dim)
Давайте перейдём к сборке блока Transformer Encoder:
class TransformerEncoderBlock(nn.Module): def __init__(self, dim, num_heads, mlp_dim, dropout): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = MultiHeadAttention(dim, num_heads, dropout) self.norm2 = nn.LayerNorm(dim) self.mlp = nn.Sequential( nn.Linear(dim, mlp_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(mlp_dim, dim), nn.Dropout(dropout) ) self.dropout = nn.Dropout(dropout) def forward(self, x, return_attn=False): if return_attn: attn_out, attn_weights = self.attn(self.norm1(x), return_attn=True) x = x + self.dropout(attn_out) x = x + self.dropout(self.mlp(self.norm2(x))) return x, attn_weights else: x = x + self.dropout(self.attn(self.norm1(x))) x = x + self.dropout(self.mlp(self.norm2(x))) return x
Здесь мы просто следуем архитектуре нашей сети — все необходимые блоки мы уже реализовали.
Мы уже почти на финишной прямой — теперь соберём сам Vision Transformer:
class VisualTransformer(nn.Module): def __init__(self,num_classes, img_size=32, patch_size=4, in_channels=3, embed_dim=256, num_layers=6, num_heads=7, mlp_dim=512, dropout=0.1): super().__init__() self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim) self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim)) self.pos_embed = nn.Parameter(torch.randn(1, 1 + self.patch_embed.num_patches, embed_dim)) self.dropout = nn.Dropout(dropout) self.encoder_blocks = nn.ModuleList([ TransformerEncoderBlock(embed_dim, num_heads, mlp_dim, dropout) for _ in range(num_layers) ]) self.norm = nn.LayerNorm(embed_dim) self.mlp_head = nn.Sequential( nn.LayerNorm(embed_dim), nn.Linear(in_features=embed_dim, out_features=num_classes) ) def forward(self, x, return_attn = False): B = x.size(0) x = self.patch_embed(x) # (B, N, D) cls_tokens = self.cls_token.expand(B, -1, -1) # (B, 1, D) x = torch.cat((cls_tokens, x), dim=1) # (B, 1+N, D) x = x + self.pos_embed x = self.dropout(x) attn_maps = [] for block in self.encoder_blocks: if return_attn: x, attn = block(x, return_attn=True) attn_maps.append(attn) # (B, heads, N, N) else: x = block(x) # (B, 1+N, D) x = self.norm(x) out = self.mlp_head(x[:, 0, :]) if return_attn: return out, attn_maps else: return out
Здесь нужно уточнить несколько моментов. Что такое cls_token? Это специальный токен, который мы добавляем вручную, и он имеет тот же размер, что и патчи изображения. Его задача — использоваться позже для классификации изображения. Идея в том, что, проходя через блоки внимания, этот токен собирает информацию обо всём изображении.
Далее посмотрим на pos_embed. Поскольку мы делим изображение на патчи и выстраиваем их в последовательность — как будто работаем с текстом — модель изначально не понимает пространственные взаимосвязи между патчами. Чтобы это исправить, мы добавляем позиционную информацию к патчам. В нашем случае pos_embed — это обучаемый параметр.
Что касается mlp_head, здесь всё просто: он берёт cls_token, пропускает его через линейный слой и классифицирует изображение.
После сборки нашей модели давайте перейдём к обучению.
Для обучения мы будем использовать следующие гиперпараметры:
BATCH_SIZE = 128 EPOCHS = 80 LEARNING_RATE = 3e-4 PATCH_SIZE = 4 NUM_CLASSES = 10 IMAGE_SIZE = 32 CHANNELS = 3 EMBED_DIM = 256 NUM_HEADS = 8 DEPTH = 6 MLP_DIM = 512 DROP_RATE = 0.1
Давайте посмотрим на количество параметров:
Total parameters: 3,189,514 Trained parameters: 3,189,514
А также следующие аугментации:
train_transforms = transforms.Compose([ transforms.Resize((70, 70)), transforms.RandomResizedCrop(32, scale=(0.8, 1.0)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(15), transforms.ColorJitter(0.2, 0.2, 0.2, 0.1), transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), transforms.RandomErasing(p=0.25, scale=(0.02, 0.2), ratio=(0.3, 3.3), value='random'), ])
Мы обучаем модель и получаем следующие результаты:
Давайте посмотрим на предсказания модели:
Вот метрики получившейся модели:
А теперь к самой интересной части — вниманию. Давайте посмотрим, на что наша модель обращает внимание во время классификации:
В этой статье мы подробно рассмотрели реализацию Vision Transformer и его механизма внимания. Мы изучили, на что способна эта модель и как она «смотрит» на изображение с помощью механизма внимания. Vision Transformer открыл новые направления в исследовании компьютерного зрения, объединив идеи из NLP и обработки изображений. В будущем мы обязательно применим эту модель для задачи генерации подписей к изображениям (Image Captioning).
Полный код и процесс обучения вы можете найти на моём Kaggle:
https://www.kaggle.com/code/nickr0ot/visual-transformer-from-scratch
