Учебник BYOL: самостоятельное обучение изображениям CIFAR с кодом в Pytorch

Искусственный интеллект

Учебник BYOL: самостоятельное обучение изображениям CIFAR с кодом в Pytorch | DeepTech

DEEPTECH

July 14, 2023

После презентации SimCLR контрастный самоконтролируемой системы обучения, я решил продемонстрировать еще один печально известный метод, называемый BYOL. Bootstrap Your Own Latent (BYOL) — это новый алгоритм самоконтролируемого обучения представлениям изображений. BYOL имеет два основных преимущества:

Он явно не использует отрицательные образцы. Вместо этого он напрямую минимизирует сходство представлений одного и того же изображения в другом расширенном представлении (положительная пара). Отрицательные образцы — это изображения из партии, отличной от положительной пары.
Как результат, BYOL утверждается, что для него требуются партии меньшего размера, что делает его привлекательным выбором.

Ниже вы можете ознакомиться с методом. В отличие от оригинальной статьи, я называю онлайн-сетевого студента и целевого сетевого учителя.

бель-обзор

Обзор метода BYOL. Источник: документ BYOL.

Онлайн-сеть ака студент: по сравнению с SimCLR есть второй MLP, называемый предсказатель, что делает весь метод асимметричным. Асимметричность по сравнению с чем? Ну и к модели учителя (целевая сеть).

Почему это важно?

Поскольку модель учителя обновлена только через экспоненциальную скользящую среднюю (EMA) от параметров студента. В конечном итоге на каждой итерации учителю передается крошечный процент (менее 1%) параметров ученика. Таким образом, градиенты проходят только через студенческую сеть. Это может быть реализовано как:

class EMA():
   def __init__(self, alpha):
       super().__init__()
       self.alpha = alpha
   def update_average(self, old, new):
       if old is None:
           return new
       return old * self.alpha + (1 - self.alpha) * new
ema = EMA(0.99)
for student_params, teacher_params in zip(student_model.parameters(),teacher_model.parameters()):
         old_weight, up_weight = teacher_params.data, student_params.data
         teacher_params.data = ema.update_average(old_weight, up_weight)

Еще одним ключевым отличием Simclr от BYOL является функция потерь.

Функция потерь

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

бель-бумага-обзор-с-тензорами

Обзор метода BYOL. Источник: документ BYOL.

Наконец, авторы определили следующую среднеквадратичную ошибку между L2-нормализованными прогнозами и целевыми прогнозами:

\mathcal{L}_{\theta, \xi}\triangleq\left\|\bar{q}_{\theta}\left(z_{\theta}\right)-\bar{z}_{\xi }^{\prime}\right\|_{2}^{2}=2-2\cdot\frac{\left\langle q_{\theta}\left(z_{\theta}\right), z_{ \xi}^{\prime}\right\rangle}{\left\|q_{\theta}\left(z_{\theta}\right)\right\|_{2} \cdot\left\|z_{ \xi}^{\prime}\right\|_{2}} .

Потеря L2 может быть реализована следующим образом. Нормализация L2 применяется заранее.

import torch
import torch.nn.functional as F
def loss_fn(x, y):
   x = F.normalize(x, dim=-1, p=2)
   y = F.normalize(y, dim=-1, p=2)
   return 2 - 2 * (x * y).sum(dim=-1)

Код доступен на Гитхаб

Отслеживание того, что происходит во время предварительной подготовки с самоконтролем: точность KNN

Тем не менее, потери в самоконтролируемом обучении не являются надежным показателем для отслеживания. Я обнаружил, что лучший способ отслеживать, что происходит во время тренировки, — это измерять точность ΚΝΝ.

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

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

import numpy as np
import torch
from sklearn.model_selection import cross_val_score
from sklearn.neighbors import KNeighborsClassifier
from torch import nn
class KNN():
   def __init__(self, model, k, device):
       super(KNN, self).__init__()
       self.k = k
       self.device = device
       self.model = model.to(device)
       self.model.eval()
   def extract_features(self, loader):
       """
       Infer/Extract features from a trained model
       Args:
           loader: train or test loader
       Returns: 3 tensors of all:  input_images, features, labels
       """
       x_lst = ()
       features = ()
       label_lst = ()
       with torch.no_grad():
           for input_tensor, label in loader:
               h = self.model(input_tensor.to(self.device))
               features.append(h)
               x_lst.append(input_tensor)
               label_lst.append(label)
           x_total = torch.stack(x_lst)
           h_total = torch.stack(features)
           label_total = torch.stack(label_lst)
           return x_total, h_total, label_total
   def knn(self, features, labels, k=1):
       """
       Evaluating knn accuracy in feature space.
       Calculates only top-1 accuracy (returns 0 for top-5)
       Args:
           features: (... , dataset_size, feat_dim)
           labels: (... , dataset_size)
           k: nearest neighbours
       Returns: train accuracy, or train and test acc
       """
       feature_dim = features.shape(-1)
       with torch.no_grad():
           features_np = features.cpu().view(-1, feature_dim).numpy()
           labels_np = labels.cpu().view(-1).numpy()
           self.cls = KNeighborsClassifier(k, metric="cosine").fit(features_np, labels_np)
           acc = self.eval(features, labels)
       return acc
   def eval(self, features, labels):
     feature_dim = features.shape(-1)
     features = features.cpu().view(-1, feature_dim).numpy()
     labels = labels.cpu().view(-1).numpy()
     acc = 100 * np.mean(cross_val_score(self.cls, features, labels))
     return acc
   def _find_best_indices(self, h_query, h_ref):
       h_query = h_query / h_query.norm(dim=1).view(-1, 1)
       h_ref = h_ref / h_ref.norm(dim=1).view(-1, 1)
       scores = torch.matmul(h_query, h_ref.t())  
       score, indices = scores.topk(1, dim=1)  
       return score, indices
   def fit(self, train_loader, test_loader=None):
       with torch.no_grad():
           x_train, h_train, l_train = self.extract_features(train_loader)
           train_acc = self.knn(h_train, l_train, k=self.k)
           if test_loader is not None:
               x_test, h_test, l_test = self.extract_features(test_loader)
               test_acc = self.eval(h_test, l_test)
               return train_acc, test_acc

Теперь мы можем сосредоточиться на методе и модели BYOL.

Изменить resnet: добавить проекционные головки MLP

Мы начнем с базовой модели (resnet18) и изменим ее для самостоятельного обучения. Последний слой, который обычно выполняет классификацию, заменяется функцией идентификации. Выходные данные resnet18 будут переданы проектору MLP.

import copy
import torch
from torch import nn
import torch.nn.functional as F
class MLP(nn.Module):
   def __init__(self, dim, embedding_size=256, hidden_size=2048, batch_norm_mlp=False):
       super().__init__()
       norm = nn.BatchNorm1d(hidden_size) if batch_norm_mlp else nn.Identity()
       self.net = nn.Sequential(
           nn.Linear(dim, hidden_size),
           norm,
           nn.ReLU(inplace=True),
           nn.Linear(hidden_size, embedding_size)
       )
   def forward(self, x):
       return self.net(x)
class AddProjHead(nn.Module):
   def __init__(self, model, in_features, layer_name, hidden_size=4096,
                embedding_size=256, batch_norm_mlp=True):
       super(AddProjHead, self).__init__()
       self.backbone = model
       setattr(self.backbone, layer_name, nn.Identity())
       self.backbone.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
       self.backbone.maxpool = torch.nn.Identity()
       self.projection = MLP(in_features, embedding_size, hidden_size=hidden_size, batch_norm_mlp=batch_norm_mlp)
   def forward(self, x, return_embedding=False):
       embedding = self.backbone(x)
       if return_embedding:
           return embedding
       return self.projection(embedding)

Я также заменил первый конверсионный слой resnet18 с свертки 7×7 на 3×3, так как мы играем с изображениями 32×32 (CIFAR-10).

Код доступен на Гитхаб. Если вы планируете закрепить свои знания о Pytorch, мы настоятельно рекомендуем две замечательные книги: Глубокое обучение с PyTorch от Manning Publications и Машинное обучение с PyTorch и Scikit-Learn Себастьян Рашка. Вы всегда можете использовать код скидки 35% блейсаммер21 для всей продукции Manning’s.

Фактический метод BYOL

До сих пор я представил все важные компоненты для достижения этой точки. Теперь будем строить BYOL модуль с нашими любимыми сетями учеников и учителей. Обратите внимание, что студенческий предиктор MLP и проектор идентичны.

Моя реализация BYOL была основана на lucidrains. репо. Я изменил его, чтобы сделать его более простым и поиграть с ним.

class BYOL(nn.Module):
   def __init__(
           self,
           net,
           batch_norm_mlp=True,
           layer_name='fc',
           in_features=512,
           projection_size=256,
           projection_hidden_size=2048,
           moving_average_decay=0.99,
           use_momentum=True):
       """
       Args:
           net: model to be trained
           batch_norm_mlp: whether to use batchnorm1d in the mlp predictor and projector
           in_features: the number features that are produced by the backbone net i.e. resnet
           projection_size: the size of the output vector of the two identical MLPs
           projection_hidden_size: the size of the hidden vector of the two identical MLPs
           augment_fn2: apply different augmentation the second view
           moving_average_decay: t hyperparameter to control the influence in the target network weight update
           use_momentum: whether to update the target network
       """
       super().__init__()
       self.net = net
       self.student_model = AddProjHead(model=net, in_features=in_features,
                                        layer_name=layer_name,
                                        embedding_size=projection_size,
                                        hidden_size=projection_hidden_size,
                                        batch_norm_mlp=batch_norm_mlp)
       self.use_momentum = use_momentum
       self.teacher_model = self._get_teacher()
       self.target_ema_updater = EMA(moving_average_decay)
       self.student_predictor = MLP(projection_size, projection_size, projection_hidden_size)
   @torch.no_grad()
   def _get_teacher(self):
       return copy.deepcopy(self.student_model)
   @torch.no_grad()
   def update_moving_average(self):
       assert self.use_momentum, 'you do not need to update the moving average, since you have turned off momentum ' \
                                 'for the target encoder '
       assert self.teacher_model is not None, 'target encoder has not been created yet'
       for student_params, teacher_params in zip(self.student_model.parameters(), self.teacher_model.parameters()):
         old_weight, up_weight = teacher_params.data, student_params.data
         teacher_params.data = self.target_ema_updater.update_average(old_weight, up_weight)
   def forward(
           self,
           image_one, image_two=None,
           return_embedding=False):
       if return_embedding or (image_two is None):
           return self.student_model(image_one, return_embedding=True)
       student_proj_one = self.student_model(image_one)
       student_proj_two = self.student_model(image_two)
       student_pred_one = self.student_predictor(student_proj_one)
       student_pred_two = self.student_predictor(student_proj_two)
       with torch.no_grad():
           teacher_proj_one = self.teacher_model(image_one).detach_()
           teacher_proj_two = self.teacher_model(image_two).detach_()
       loss_one = loss_fn(student_pred_one, teacher_proj_one)
       loss_two = loss_fn(student_pred_two, teacher_proj_two)       
       return (loss_one + loss_two).mean()

Для CIFAR-10 достаточно использовать 2048 в качестве скрытого измерения и 256 в качестве встраиваемого измерения. Мы будем обучать resnet18, который выводит 512 признаков за 100 эпох. Части кода, относящиеся к загрузке и дополнению данных, опущены для удобства чтения. Вы можете посмотреть их в коде.

Вы можете использовать оптимизатор Адама ( $лр=3 * 10^{-4}$

Единственное, что будет изменено в коде поезда, — это обновление EMA.

def training_step(model, data):
   (view1, view2), _ = data
   loss = model(view1.cuda(), view2.cuda())
   return loss
def train_one_epoch(model, train_dataloader, optimizer):
   model.train()
   total_loss = 0.
   num_batches = len(train_dataloader)
   for data in train_dataloader:
       optimizer.zero_grad()
       loss = training_step(model, data)
       loss.backward()
       optimizer.step()
       model.update_moving_average()
       total_loss += loss.item()
   return total_loss/num_batches

Давайте прыгать на результаты!

Результаты: точность KNN по сравнению с эпохами предварительной подготовки

knn-byol-обучение

Точность KNN каждые 4 эпохи. Изображение автора

Разве не удивительно, что без каких-либо меток мы можем достичь точности валидации в 70%? Я нашел это удивительным, особенно для этого метода, который кажется менее чувствительным к размеру пакета.

Но почему здесь влияет размер партии? Разве это не должно не использовать отрицательный париж? Откуда берется зависимость размера партии?

Краткий ответ: Ну, это нормализация пакетов в слоях MLP!

Вот эксперименты, которые я провел, чтобы проверить это.

Примечание о пакетной норме в сетях MLP и моментуме EMA

Мне было любопытно наблюдать за крахом режима без пакетной нормализации. Вы можете попробовать это самостоятельно, установив:

model = BYOL(model, in_features=512, batch_norm_mlp=False)

Я заметил, что расстояние L2 стремится почти к нулю с самых первых эпох:

Epoch 0: loss:0.06423207696957084

Epoch 8: loss:0.005584242034894534

Epoch 20: loss:0.005460431350347323

Убыток приближается к нулю, и KNN перестает расти (35% против 60% в обычном режиме). Вот почему утверждается, что BYOL неявно использует форму сравнительного обучения, используя статистику партии в MLP. Вот точность KNN:

Свертывание режима в BYOL путем удаления пакетной нормы в MLP. Изображение автора

Мне хорошо известны документы, показывающие, что пакетная статистика — не единственное условие для работы BYOL. Это экспериментальный пост, поэтому я не буду играть в эту игру. Мне просто было любопытно понаблюдать здесь за коллапсом режима.

Заключение

Для более подробного объяснения метода посмотрите видео Янника на BYOL:

В этом руководстве мы шаг за шагом внедрили BYOL и предварительно обучили CIFAR10. Мы наблюдаем значительное увеличение точности KNN за счет сопоставления представлений одного и того же изображения. Случайный классификатор будет иметь 10%, а со 100 эпохами мы достигнем 70% точности проверки KNN без каких-либо меток. Как это круто?

Чтобы узнать больше об обучении с самоконтролем, следите за обновлениями! Поддержите нас, поделившись в социальных сетях, сделав пожертвование или купив нашу книгу «Глубокое обучение в производстве». Было бы очень признательно.

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