Изображение стоит 16x16 слов: трансформеры для распознавания изображений в масштабе

Это документ из исследования Google. Основная идея этой статьи состоит в том,

Преобразователи, применяемые непосредственно к участкам изображения и предварительно обученные на больших наборах данных, очень хорошо работают с классификацией изображений.

В этом посте мы подробно обсудим архитектуру Vision Transformers (ViT) и результаты, опубликованные в статье.

Попробуйте реализации, использованные здесь, с помощью этого блокнота. Нажмите Open in Colab, чтобы запустить прямо в Colab.

Архитектура ВиТ 🏃

ViT разбивает изображение на фиксированное количество патчей и использует их для создания вложений и пропускает их через стандартный кодер-трансформер.
Давайте обсудим, как работает преобразователь НЛП, а затем сравним и обсудим его с ViT.

Кодер-трансформер НЛП 🎨

Мы берем предложение в качестве входных данных. Но вместо отправки всего предложения мы используем tokenizer, что дает каждому слову id. Теперь токенизатору на самом деле не нужно разбивать наивным образом, как разбивать по словам, но он может разбивать слово на части и также присваивать. Это зависит от того, как обучен токенизатор. Например, простой токенизатор может сделать ниже

Но другой разделил предложение, как показано ниже:

Теперь у нас есть матрица, которая отображает все возможные идентификаторы в векторное представление. Итак, если мы используем второй токенизатор, если мы хотим выбрать встраивание для To, мы выберем векторное представление по индексу 11.

Эти вложения могут быть случайным образом инициализированы в начале, и мы учимся во время обучения.

Позиционные вложения

Рекуррентные нейронные сети (RNN) последовательно разбирают предложение слово за словом. Но архитектура Transformer не использует механизм повторения в пользу многоголового механизма самоконтроля. Это сокращает время обучения в преобразователях, но модель не имеет представления о положении слов.
Чтобы решить эту проблему, мы добавляем дополнительную информацию (позиционные кодировки) к входным вложениям.
Теперь один простой способ — просто присвоить 1 первому слову, 2 — второму и так далее. Но при таком подходе модель во время логического вывода может получить более длинное предложение, чем любое, которое она видела во время обучения. Кроме того, для более длинного предложения будут добавлены большие значения, которые занимают больше памяти.
Мы можем взять диапазон, а затем добавить 0 для первой работы и 1 для последней, все, что между ними, мы разделяем диапазон [0,1] и получаем значения. Например, для предложения из 3 слов мы можем сделать 0 для первого слова, 0,5 для второго и 1 для третьего; для предложения из 4 слов это будет 0,0,33, 0,66, 1 соответственно. Проблема в том, что дельта разницы позиций не является постоянной. В первом примере это было 0,5, а во втором — 0,33.
Используемое позиционное кодирование представляет собой d-мерный вектор.

Так вот как все совпало

Далее мы пропустим эти векторы через многоголовые блоки внимания.

Многоголовое внимание 🔥

Многоголовое внимание имеет три матрицы: запрос (Q), ключ (K) и матрицу значения (V). Каждая из них имеет те же размеры, что и закладная. Итак, в нашем случае все 3 матрицы 512х512
Для каждого вложения токена мы умножаем его на все три матрицы (Q, K, V). Таким образом, у нас будет 3 промежуточных вектора длиной 512 для каждого токена.
Теперь, если у нас есть n головок, мы делим каждый из этих векторов на n частей. Например, если у нас 8 голов, то для слова Today мы разделим все 3 промежуточных вектора на маленькие векторы размерностью 64.
Затем каждая головка берет соответствующий ей сегмент из всех промежуточных векторов. Например, первая головка возьмет первое разделение (размерность 64) всех трех промежуточных векторов (соответствующих результатам умножения запроса, ключа, значения) всех пяти вложений (соответствующих пяти токенов). Точно так же вторая головка займет второй сегмент и так далее.
В каждой голове мы производим точечный продукт между векторами, умноженными на запрос и ключевую матрицу. На изображении ниже для головы 1 мы делаем скалярное произведение между q1 и всем вектором, умноженным на ключевую матрицу (k{i}, i в [1,5]). Затем мы умножаем его на соответствующий вектор значений. Наконец, мы добавляем их, чтобы создать результирующий 64-мерный вектор. Это происходит для q2, q3, q4, q5 и, наконец, мы получаем 5 векторов с размерностью 64. Теперь в основном каждый результирующий вектор содержит информацию обо всех других векторах.

Теперь мы объединяем результирующие векторы со всех головок. Таким образом, мы соединим первый результирующий вектор из всех 8 головок, чтобы создать первый 512-мерный вектор. То же самое происходит для всех остальных 4 векторов.
Итак, наконец, у нас есть 5 векторов по 512 измерений каждый.

Добавить и нормировать ☯

Это обычная пакетная нормализация и остаточные соединения, такие как блок Resnet.

Прямая связь 🍀

Это простая нейронная сеть с прямой связью, которая применяется к каждому вектору внимания.

Вот как работает простой энкодер трансформаторов. Давайте посмотрим на архитектуру ViT дальше. Мы также увидим, как это реализовать одновременно.

Архитектура энкодера ViT

Встроенные исправления ☑️

Для обработки 2D-изображений изображение делится на несколько фрагментов. и мы сглаживаем эти 2D-патчи до 1D-векторов.
Затем мы вкладываем каждый из этих векторов в размерное пространство модели. В этом случае модель преобразует каждый вектор в 768 векторов измерений.

import torch
import torch.nn as nn
in_chans = 3 #RGB
embed_dim = 768 # vector dimension in model space
patch_size = 16 # each image patch size 16*16
proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) # this will create the patch in image
img = torch.randn(1, 3, 224,224) # dummy image
x = proj(img).flatten(2).transpose(1, 2) # BCHW -> BNC
print(x.shape)

В приведенном выше коде мы взяли изображения размером 224*224 и предположили, что каждый патч имеет размер 16x16.
Теперь это даст в сумме (224/16 * 224/16 ) = 14*14 = 196 векторов.
Каждый из этих векторов размером 16 * 16 = 256. Но, поскольку мы должны преобразовать его в размерность модели, которая равна 768, мы используем 768 в качестве выходных каналов в свертке. Наконец, мы сглаживаем его до BNC, где B = пакет, N = результирующие патчи, C = векторное измерение в пространстве модели.

Встраивания классов 🆕

ViT добавляет обучаемое встраивание к последовательности встроенных патчей.

cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) # create class embeddings without batch
cls_token = cls_token.expand(x.shape[0], -1, -1) # add batch
x = torch.cat((cls_token, x), dim=1) # append class token with linear proj embeddings
x.shape # 196 -> 197

позиционные вложения ☑️

Создаем матрицу размерности (num_patches+1) * embed_dim (197*768). Значения изучаются во время обучения.

num_patches = 14*14
pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) # +1 for class token
x = x + pos_embed # add position encoding
x.shape

Блокирует ☑️

На основе модели у нас есть n количество блоков.
Каждый блок одинаковый. Каждый из них состоит из уровня внимания и уровня MLP.

Слой внимания ☑️

То же, что и внимание в НЛП, которое объяснялось ранее.
Создадим промежуточные векторы.

# Transformation from source vector to query vector
fc_q = nn.Linear(embed_dim, embed_dim)
# Transformation from source vector to key vector
fc_k = nn.Linear(embed_dim, embed_dim)
# Transformation from source vector to value vector
fc_v = nn.Linear(embed_dim, embed_dim)
Q = fc_q(x)
K = fc_k(x)
V = fc_v(x)
print(Q.shape, K.shape, V.shape)

Разделите промежуточные векторы, чтобы обработать одну часть в каждой головке.

num_heads = 8
batch_size = 1
Q = Q.view(batch_size, -1, num_heads, embed_dim//num_heads).permute(0, 2, 1, 3) # split the Q matrix for 8 head
K = K.view(batch_size, -1, num_heads, embed_dim//num_heads).permute(0, 2, 1, 3) # split the K matrix for 8 head
V = V.view(batch_size, -1, num_heads, embed_dim//num_heads).permute(0, 2, 1, 3) # split the V matrix for 8 head
print(Q.shape, K.shape, V.shape) # batch_size, num_head, num_patch+1, feature_vec dim per head

Умножение матриц внимания.

score = torch.matmul(Q, K.permute(0, 1, 3, 2)) # Q*k
score = torch.softmax(score, dim=-1)
score = torch.matmul(score, V) # normally we apply dropout layer before this
score.shape # batch_size, num_head, num_patches+1, feature_vector_per_head (embed_dim/num_head)

Изменить результаты

score = score.permute(0, 2, 1, 3).contiguous()
score.shape # batch_size, num_patches+1, num_head, feature_vector_per_head (embed_dim/num_head)

Объедините векторы обратно в их первоначальную форму

score = score.view(batch_size, -1, embed_dim) # merge the vectors back to original shape
score.shape # batch_size, num_patches+1, embed_dim

Глава MLP ☑️

Обычный многослойный персептрон.

act_layer=nn.GELU # activation function
in_features = embed_dim 
hidden_features = embed_dim * 4
out_features = in_features
fc1 = nn.Linear(in_features, hidden_features)
act = act_layer()
drop1 = nn.Dropout(0.5)
fc2 = nn.Linear(hidden_features, out_features)
drop2 = nn.Dropout(0.5)

Получите результат от слоев MLP

x = fc1(score)
x = act(x)
x = drop1(x)
x = fc2(x)
x = drop2(x)
x.shape

выньте функции токена cls.

cls = x[:,0]

Начало классификатора 🆕

Создайте простой заголовок классификатора и передайте функции токена класса, чтобы получить прогнозы.

num_classes = 10 # assume 10 class classification
head = nn.Linear(embed_dim, num_classes) 
pred = head(cls)
pred

Результаты опубликованы в статье 📈

При обучении на наборах данных среднего размера, таких как ImageNet, без строгой регуляризации эти модели дают скромную точность на несколько процентных пунктов ниже, чем у ResNet сопоставимого размера.

Преобразователям не хватает некоторых индуктивных смещений, присущих CNN, таких как эквивалентность перевода и локальность, и поэтому они плохо обобщают при обучении на недостаточном количестве данных.

Однако картина меняется, если модели обучаются на больших наборах данных (14–300 млн изображений). Мы обнаружили, что крупномасштабное обучение превосходит индуктивную предвзятость.

Авторы отметили, что для небольших наборов данных для предварительной подготовки (ImageNet) модели ViT-Large уступают моделям ViT-Base. С большими наборами данных (JFT-300M) хорошо работают модели ViT-Large.
Модели Vision Transformer, предварительно обученные на наборе данных JFT-300M, превосходят базовые модели на основе ResNet для всех наборов данных, при этом для предварительной подготовки требуется значительно меньше вычислительных ресурсов.
В следующей таблице показан результат предварительной подготовки ViT с набором данных JFT-300M и набором данных ImageNet-21k. Столбцы показывают несколько моделей, предварительно обученных с разными наборами данных. Строки — это последующие задачи.

Тренируйте простой ViT с PyTorch Lightning и timm 🎆

Здесь давайте обучим простой классификатор, используя PyTorch Lightning и timm.

import timm
import torch
import pytorch_lightning as pl
import torchvision
import torchvision.transforms as transforms
from pytorch_lightning import Trainer, seed_everything
from pytorch_lightning.callbacks import ModelCheckpoint
import torchmetrics
seed_everything(42, workers=True)

Давайте создадим простой класс модели молнии.

class Model(pl.LightningModule):
    """
    Lightning model
    """
    def __init__(self, model_name, num_classes, lr = 0.001, max_iter=20):
        super().__init__()
        self.model = timm.create_model(model_name=model_name, pretrained=True, num_classes=num_classes)
        self.metric = torchmetrics.Accuracy()
        self.loss = torch.nn.CrossEntropyLoss()
        self.lr = lr
        self.max_iter = max_iter
        
    def forward(self, x):
        return self.model(x)
    def shared_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        self.metric(preds, y)
        
        return loss
    
    def training_step(self, batch, batch_idx):
        loss = self.shared_step(batch, batch_idx)
        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, prog_bar=True)
        self.log('train_acc', self.metric, on_epoch=True, logger=True, prog_bar=True)
        
        return loss
    
    def validation_step(self, batch, batch_idx):
        loss = self.shared_step(batch, batch_idx)
        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, prog_bar=True)
        self.log('val_acc', self.metric, on_epoch=True, logger=True, prog_bar=True)
        
        return loss
    
    def configure_optimizers(self):
        optim = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer=optim, T_max=self.max_iter)
        
        return [optim], [scheduler]

Далее мы определим преобразования, а также загрузим и загрузим набор данных CIFAR10.

transform = transforms.Compose(
    [transforms.Resize(224),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 128
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=8)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                         shuffle=False, num_workers=8)
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

Теперь мы инициализируем класс Model. Здесь мы используем вариант ViT, который берет изображения размером 224*224, а размер патча равен 16.

model = Model(model_name="vit_tiny_patch16_224", num_classes=len(classes), lr = 0.001, max_iter=10)

Давайте создадим обратный вызов контрольной точки, чтобы сохранить лучшую контрольную точку.

checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    dirpath='./checkpoints',
    filename='vit_tpytorch_lightning6_224-cifar10-{epoch:02d}-{val_loss:.2f}-{val_acc:.2f}'
)

Почти готово. Создадим трейнер.

trainer = Trainer(
    deterministic=True, 
    logger=False, 
    callbacks=[checkpoint_callback], 
    gpus=[0], # change it based on gpu or cpu availability
    max_epochs=10, 
    stochastic_weight_avg=True)

Наконец, давайте обучим модель 😃

trainer.fit(model=model, train_dataloaders=trainloader, val_dataloaders=testloader)

Связанные ресурсы

Особая благодарность Пракашу Джею за помощь в этом проекте.

Предложения по подаче заявок на Mlearning.ai
Как стать писателем на Mlearning.aimedium.com