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

В предыдущем уроке я немного рассказал об обучении с самостоятельным наблюдением. Пришло время приступить к вашему первому проекту, запустив SimCLR на небольшом наборе данных со 100 000 немаркированных изображений под названием СТЛ10.

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

Метод SimCLR: контрастное обучение

Позволять $сим (и, в)$ обратите внимание на скалярное произведение между 2 нормализованными $ты$ и $в$ векторов (т.е. косинусное сходство).

Тогда функция потерь для положительной пары примеров (i,j) определяется как:

где $\mathbb{1}_{(k \neq i)} \in {0,1}$

$\тау$ обозначает температурный параметр. Окончательный убыток вычисляется путем суммирования всех положительных пар и деления на $2\x N = просмотры \time пакет\_size$

Существуют различные способы развития контрастных потерь. Здесь мы предоставляем вам важную информацию.

Нормализация L2 и вычисление матрицы подобия косинусов

Во-первых, к признакам нужно применить нормализацию L2, иначе этот метод не работает. Нормализация L2 означает, что векторы нормированы таким образом, что все они лежат на поверхности единичной (гипер)сферы, где норма L2 равна 1.

z_i = F.normalize(proj_1, p=2, dim=1)
z_j = F.normalize(proj_2, p=2, dim=1)

Объедините 2 выходных представления в пакетном измерении. Их форма будет $(2 \раза партия\_размер, тусклый)$

def calc_similarity_batch(self, a, b):
       representations = torch.cat((a, b), dim=0)
       return F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)

Индексирование матрицы сходства для функции потерь SimCLR

Теперь нам нужно проиндексировать полученную матрицу размера $(пакет\_размер\количество просмотров, пакет\_размер\количество просмотров)$

Наглядная иллюстрация SimCLR. Изображение от автора

Хорошо, как, черт возьми, мы это делаем? У меня такой же вопрос. Здесь размер пакета составляет 2 изображения, но мы хотим реализовать решение для любого размера пакета. Если вы внимательно посмотрите, то увидите, что положительные пары сдвинуты от главной диагонали на 2, то есть на размер партии. Один из способов сделать это torch.diag(). Он берет выбранную диагональ из матрицы. Первый параметр — это матрица, а второй указывает диагональ, где нуль представляет элементы главной диагонали. Берем диагонали, сдвинутые на размер партии.

sim_ij = torch.diag(similarity_matrix, batch_size)
sim_ji = torch.diag(similarity_matrix, -batch_size)
positives = torch.cat((sim_ij, sim_ji), dim=0)

Есть $пакет\_размер\количество просмотров$

(0., 0., 0., 1., 0., 0.),
(0., 0., 0., 0., 1., 0.),
(0., 0., 0., 0., 0., 1.),
(1., 0., 0., 0., 0., 0.),
(0., 1., 0., 0., 0., 0.),
(0., 0., 1., 0., 0., 0.)

Для знаменателя нам нужны как положительные, так и отрицательные пары. Таким образом, бинарная маска будет точной поэлементной инверсией единичной матрицы.

self.mask = (~torch.eye(batch_size * 2, batch_size * 2, dtype=bool)).float()
pos_and_negatives = self.mask * similarity_matrix

Опять же, они являются как положительными, так и отрицательными в знаменателе.

Вы можете разобрать остальное (масштабирование температуры и суммирование отрицательных значений из знаменателя и т. д.):

Реализация потерь SimCLR

import torch
import torch.nn as nn
import torch.nn.functional as F
def device_as(t1, t2):
   """
   Moves t1 to the device of t2
   """
   return t1.to(t2.device)
class ContrastiveLoss(nn.Module):
   """
   Vanilla Contrastive loss, also called InfoNceLoss as in SimCLR paper
   """
   def __init__(self, batch_size, temperature=0.5):
       super().__init__()
       self.batch_size = batch_size
       self.temperature = temperature
       self.mask = (~torch.eye(batch_size * 2, batch_size * 2, dtype=bool)).float()
   def calc_similarity_batch(self, a, b):
       representations = torch.cat((a, b), dim=0)
       return F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)
   def forward(self, proj_1, proj_2):
       """
       proj_1 and proj_2 are batched embeddings (batch, embedding_dim)
       where corresponding indices are pairs
       z_i, z_j in the SimCLR paper
       """
       batch_size = proj_1.shape(0)
       z_i = F.normalize(proj_1, p=2, dim=1)
       z_j = F.normalize(proj_2, p=2, dim=1)
       similarity_matrix = self.calc_similarity_batch(z_i, z_j)
       sim_ij = torch.diag(similarity_matrix, batch_size)
       sim_ji = torch.diag(similarity_matrix, -batch_size)
       positives = torch.cat((sim_ij, sim_ji), dim=0)
       nominator = torch.exp(positives / self.temperature)
       denominator = device_as(self.mask, similarity_matrix) * torch.exp(similarity_matrix / self.temperature)
       all_losses = -torch.log(nominator / torch.sum(denominator, dim=1))
       loss = torch.sum(all_losses) / (2 * self.batch_size)
       return loss

Аугментации

Ключом к обучению репрезентации с самостоятельным наблюдением является увеличение данных. Обычно используется следующий конвейер преобразования:

• Обрезка в случайном масштабе от 7% до 100% изображения

• Измените размер всех изображений до 224 или других пространственных размеров.

• Применить горизонтальное отражение с вероятностью 50%

• Применить сильное цветовое дрожание с вероятностью 80%

• Применить размытие по Гауссу с вероятностью 50%. Размер ядра обычно составляет около 10% изображения или меньше.

• Преобразование изображений RGB в оттенки серого с вероятностью 20%.

• Нормализация на основе средних и отклонений imagenet

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

import torch
import torchvision.transforms as T
class Augment:
   """
   A stochastic data augmentation module
   Transforms any given data example randomly
   resulting in two correlated views of the same example,
   denoted x ̃i and x ̃j, which we consider as a positive pair.
   """
   def __init__(self, img_size, s=1):
       color_jitter = T.ColorJitter(
           0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s
       )
       blur = T.GaussianBlur((3, 3), (0.1, 2.0))
       self.train_transform = torch.nn.Sequential(
           T.RandomResizedCrop(size=img_size),
           T.RandomHorizontalFlip(p=0.5),  
           T.RandomApply((color_jitter), p=0.8),
           T.RandomApply((blur), p=0.5),
           T.RandomGrayscale(p=0.2),
           T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
       )
   def __call__(self, x):
       return self.train_transform(x), self.train_transform(x)

Ниже приведены 4 разных вида одного и того же изображения с применением одного и того же стохастического пайплайна:

4 разных дополнения одного и того же пайплайна. Изображение автора

Чтобы визуализировать их, вам нужно отменить нормализацию среднего стандартного значения и поместить цветовые каналы в последнее измерение:

def imshow(img):
   """
   shows an imagenet-normalized image on the screen
   """
   mean = torch.tensor((0.485, 0.456, 0.406), dtype=torch.float32)
   std = torch.tensor((0.229, 0.224, 0.225), dtype=torch.float32)
   unnormalize = T.Normalize((-mean / std).tolist(), (1.0 / std).tolist())
   npimg = unnormalize(img).numpy()
   plt.imshow(np.transpose(npimg, (1, 2, 0)))
   plt.show()
dataset = STL10("./", split='train', transform=Augment(96), download=True)
imshow(dataset(99)(0)(0))
imshow(dataset(99)(0)(0))
imshow(dataset(99)(0)(0))
imshow(dataset(99)(0)(0))

Измените Resnet18 и определите группы параметров

Одним из важных шагов для запуска simclr является удаление последнего полносвязного слоя. Мы заменим его функцией тождества. Затем нам нужно добавить проекционную головку (еще один MLP), которая будет использоваться только на этапе предварительной подготовки с самоконтролем. Для этого нам нужно знать размерность функций нашей модели. В частности, resnet18 выводит 512-мерный вектор, а resnet50 выводит 2048-мерный вектор. Проекция MLP преобразует его в размер вектора встраивания, который составляет 128, согласно официальной статье.

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

import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
from torch.optim import SGD, Adam
class AddProjection(nn.Module):
   def __init__(self, config, model=None, mlp_dim=512):
       super(AddProjection, self).__init__()
       embedding_size = config.embedding_size
       self.backbone = default(model, models.resnet18(pretrained=False, num_classes=config.embedding_size))
       mlp_dim = default(mlp_dim, self.backbone.fc.in_features)
       print('Dim MLP input:',mlp_dim)
       self.backbone.fc = nn.Identity()
       self.projection = nn.Sequential(
           nn.Linear(in_features=mlp_dim, out_features=mlp_dim),
           nn.BatchNorm1d(mlp_dim),
           nn.ReLU(),
           nn.Linear(in_features=mlp_dim, out_features=embedding_size),
           nn.BatchNorm1d(embedding_size),
       )
   def forward(self, x, return_embedding=False):
       embedding = self.backbone(x)
       if return_embedding:
           return embedding
       return self.projection(embedding)

Следующим шагом является разделение параметров моделей на 2 группы.

Целью второй группы является удаление потери веса из слоев нормализации партии. В случае использования оптимизатора LARS вам также необходимо удалить уменьшение веса из смещений. Одним из способов достижения этого является следующая функция:

def define_param_groups(model, weight_decay, optimizer_name):
   def exclude_from_wd_and_adaptation(name):
       if 'bn' in name:
           return True
       if optimizer_name == 'lars' and 'bias' in name:
           return True
   param_groups = (
       {
           'params': (p for name, p in model.named_parameters() if not exclude_from_wd_and_adaptation(name)),
           'weight_decay': weight_decay,
           'layer_adaptation': True,
       },
       {
           'params': (p for name, p in model.named_parameters() if exclude_from_wd_and_adaptation(name)),
           'weight_decay': 0.,
           'layer_adaptation': False,
       },
   )
   return param_groups

Я не использую оптимизатор LARS в этом руководстве, но если вы планируете его использовать, вот реализация который я использую в качестве ссылки.

Логика обучения SimCLR

Здесь мы реализуем всю обучающую логику SimCLR. Возьмите 2 представления, перешлите их, чтобы получить прогнозы встраивания, и рассчитайте потери SimCLR.

Мы можем завершить обучение SimCLR одним классом, используя Питорч молния который инкапсулирует всю логику обучения. В самом простом виде нам нужно реализовать training_step метод, который получает на вход пакет от загрузчика данных. Вы можете думать об этом как о вызове batch = next(iter(dataloader)) на каждом шагу. Далее идет configure_optimizers метод, связывающий модель с оптимизатором и планировщиком обучения. Я использовал уже реализованный планировщик от молнии PyTorch болты (еще один небольшой пакет в экосистеме молнии). По сути, мы постепенно увеличиваем скорость обучения до ее базового значения, а затем выполняем косинусный отжиг.

class SimCLR_pl(pl.LightningModule):
   def __init__(self, config, model=None, feat_dim=512):
       super().__init__()
       self.config = config
       self.augment = Augment(config.img_size)
       self.model = AddProjection(config, model=model, mlp_dim=feat_dim)
       self.loss = ContrastiveLoss(config.batch_size, temperature=self.config.temperature)
   def forward(self, X):
       return self.model(X)
   def training_step(self, batch, batch_idx):
       x, labels = batch
       x1, x2 = self.augment(x)
       z1 = self.model(x1)
       z2 = self.model(x2)
       loss = self.loss(z1, z2)
       self.log('Contrastive loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
       return loss
   def configure_optimizers(self):
       max_epochs = int(self.config.epochs)
       param_groups = define_param_groups(self.model, self.config.weight_decay, 'adam')
       lr = self.config.lr
       optimizer = Adam(param_groups, lr=lr, weight_decay=self.config.weight_decay)
       print(f'Optimizer Adam, '
             f'Learning Rate {lr}, '
             f'Effective batch size {self.config.batch_size * self.config.gradient_accumulation_steps}')
       scheduler_warmup = LinearWarmupCosineAnnealingLR(optimizer, warmup_epochs=10, max_epochs=max_epochs,
                                                        warmup_start_lr=0.0)
       return (optimizer), (scheduler_warmup)

Накопление градиента и эффективный размер партии

Здесь важно подчеркнуть важность использования большого размера партии. Этот метод сильно зависит от большого размера пакета, чтобы оттолкнуть от двух просмотров одного и того же изображения (позитивов). Для этого при ограниченном бюджете мы можем использовать накопление градиента. Усредняем градиенты $Н$ шаги, а затем обновлять модель вместо обновления после каждого прохода вперед-назад.

Таким образом, теперь должно быть совершенно понятно, что эффективная партия: $пакет\_размер\_на\_gpu * накопление\_steps * количество\_of\_gpu$

«В компьютерном программировании перезвонить — это ссылка на исполняемый код или фрагмент исполняемого кода, который передается в качестве аргумента другому коду. Это позволяет программному уровню более низкого уровня вызывать подпрограмму (или функцию), определенную на уровне более высокого уровня». ~ Переполнение стека

from pytorch_lightning.callbacks import GradientAccumulationScheduler
accumulator = GradientAccumulationScheduler(scheduling={0: train_config.gradient_accumulation_steps})

Основной сценарий предварительной подготовки SimCLR

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

import torch
from pytorch_lightning import Trainer
import os
from pytorch_lightning.callbacks import GradientAccumulationScheduler
from pytorch_lightning.callbacks import ModelCheckpoint
from torchvision.models import  resnet18
available_gpus = len((torch.cuda.device(i) for i in range(torch.cuda.device_count())))
save_model_path = os.path.join(os.getcwd(), "saved_models/")
print('available_gpus:',available_gpus)
filename='SimCLR_ResNet18_adam_'
resume_from_checkpoint = False
train_config = Hparams()
reproducibility(train_config)
save_name = filename + '.ckpt'
model = SimCLR_pl(train_config, model=resnet18(pretrained=False), feat_dim=512)
data_loader = get_stl_dataloader(train_config.batch_size)
accumulator = GradientAccumulationScheduler(scheduling={0: train_config.gradient_accumulation_steps})
checkpoint_callback = ModelCheckpoint(filename=filename, dirpath=save_model_path,every_n_val_epochs=2,
                                       save_last=True, save_top_k=2,monitor='Contrastive loss_epoch',mode='min')
if resume_from_checkpoint:
    trainer = Trainer(callbacks=(accumulator, checkpoint_callback),
        gpus=available_gpus,
        max_epochs=train_config.epochs,
        resume_from_checkpoint=train_config.checkpoint_path)
else:
    trainer = Trainer(callbacks=(accumulator, checkpoint_callback),
        gpus=available_gpus,
        max_epochs=train_config.epochs)
trainer.fit(model, data_loader)
trainer.save_checkpoint(save_name)
from google.colab import files
files.download(save_name)

Тонкая настройка

Хорошо, мы обучили модель. Теперь пришло время тонкой настройки. Мы будем использовать класс модуля молнии PyTorch для инкапсуляции логики. Я беру предварительно обученную основу resnet18 без проекционной головки и добавляю только один линейный слой сверху. Я точно настраиваю всю сеть. Никаких дополнений здесь не применяется. Они только отсрочат обучение. Вместо этого мы хотели бы количественно оценить производительность по сравнению с предварительно обученными весами в сети изображений и случайной инициализацией.

import pytorch_lightning as pl
import torch
from torch.optim import SGD
class SimCLR_eval(pl.LightningModule):
   def __init__(self, lr, model=None, linear_eval=False):
       super().__init__()
       self.lr = lr
       self.linear_eval = linear_eval
       if self.linear_eval:
            model.eval()
       self.mlp = torch.nn.Sequential(
            torch.nn.Linear(512,10), 
       )
       self.model = torch.nn.Sequential(
            model, self.mlp
       )
       self.loss = torch.nn.CrossEntropyLoss()
   def forward(self, X):
       return self.model(X)
   def training_step(self, batch, batch_idx):
       x, y = batch
       z = self.forward(x)
       loss = self.loss(z, y)
       self.log('Cross Entropy loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
       predicted = z.argmax(1)
       acc = (predicted == y).sum().item() / y.size(0)
       self.log('Train Acc', acc, on_step=False, on_epoch=True, prog_bar=True, logger=True)
       return loss
   def validation_step(self, batch, batch_idx):
       x, y = batch
       z = self.forward(x)
       loss = self.loss(z, y)
       self.log('Val CE loss', loss, on_step=True, on_epoch=True, prog_bar=False, logger=True)
       predicted = z.argmax(1)
       acc = (predicted == y).sum().item() / y.size(0)
       self.log('Val Accuracy', acc, on_step=True, on_epoch=True, prog_bar=True, logger=True)
       return loss
   def configure_optimizers(self):
       if self.linear_eval:
           print(f"\n\n Attention! Linear evaluation \n")
           optimizer = SGD(self.mlp.parameters(), lr=self.lr, momentum=0.9)
       else:
           optimizer = SGD(self.model.parameters(), lr=self.lr, momentum=0.9)
       return (optimizer)

Важно отметить, что STL10 является подмножество imagenet, поэтому ожидается, что трансферное обучение от imagenet будет работать очень хорошо.

Во всех случаях модель переобучается во время тонкой настройки. Помните, что никакие аугментации не применялись.

Заключение

Даже с несправедливой оценкой по сравнению с предварительно обученными весами из imagenet, контрастное самоконтролируемое обучение демонстрирует некоторые сверхмногообещающие результаты. Есть много других самоконтролируемых методов, но SimCLR является базовым.

В заключение мы рассмотрели, как шаг за шагом создать функцию потерь SimCLR и запустить обучающий скрипт без большого количества стандартного кода с помощью Pytorch-lightning. Несмотря на то, что существует разрыв между представлениями, изученными SimCLR, новейшие современные методы догоняют и даже превосходят функции, полученные с помощью imagenet, во многих областях.

Спасибо за ваш интерес к ИИ и оставайтесь позитивными!

Книга «Глубокое обучение в производстве» 📖

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

Узнать больше

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