Дообучаем готовую нейросеть для классификации данных

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

• approved — картинки идут в раздел collective (развлекательный контент и мемы);

• not suitable — не попадают в общую ленту, но остаются в ленте пользователя (селфи, пейзажи и другие);

• risked — получают бан и удаляются из приложения (расизм, порнография, расчленёнка и всё, что попадает под определение «противоправный контент»).

Сегодня на наглядных примерах расскажу, как мы перестраивали модель под наши классы, обучали её и выделяли паттерны распознавания картинок. Технические подробности — под катом.

Для начала возьмём небольшую сеть VGG-11, которая уже реализована во фреймворке pytorch. Сейчас есть много других сетей с результатами получше, но данная модель достаточно легкая, чтобы получить заметный результат за короткое время. 

Зададим несколько преобразований, чтобы унифицировать данные и привести их к привычному для сети виду. Из документации следует, что модель была предобучена на изображениях размера 224×224 пикселя, приведённых к такому формату с помощью преобразований:

from torchvision import transforms

transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406], 
            std=[0.229, 0.224, 0.225]
        )
    ]
)

Сначала уменьшаем изображение вдоль его наименьшей стороны до 256 пикселей, а затем вырезаем из центра квадрат со сторонами 224×224 пикселя. После этого независимо производится нормировка вдоль каждого канала в RGB-пространстве. Именно поэтому в преобразовании задано три значения среднего, и столько же дисперсий.

Зачем вообще нужно делать нормализацию? Представьте, что вы непутёвая Золушка и рассыпали на пол несколько видов круп, которые нужно собрать и разделить. Можно собирать по зёрнышку и сразу откладывать в нужную кучку, но удобнее собрать всё в одну кучку, а уже потом разделять по сортам. То же самое мы делаем для сети — собираем все значения в одну кучу, а потом заставляем сеть делить их на выделенные классы. 

Теперь загрузим саму модель и её веса, полученные при обучении на датасете ImageNet:

From torchvision import models
model = models.vgg11(pretrained=True).eval()

Метод eval класса VGG позволяет отключить обучение и расчёт градиента, что ускоряет предсказание модели там, где нужно узнать только ответ без дополнительного обучения, а также фиксирует все веса, что позволяет получать один и тот же ответ независимо от количества запусков. Картинки, которые находятся в наборе данных ImageNet, выглядят примерно так:

Посмотрим, на какие объекты изображений реагирует сеть при выборе того или иного класса — так лучше поймём, как она работает. Для этого воспользуемся вектором Шепли (Shapley Value), который отражает то, как меняется решение сети с параметром и без него (в нашем случае параметром будет значение пикселя).

Код для получения рисунка ниже можно взять из документации. Мы лишь добавили ещё одно изображение, заменили модель (в примере используется VGG-16), а также слой, с которого берётся градиент (с седьмого, как в примере, на десятый).

Результат следующий:

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

Если погуглить названия классов, то будет видно, что даже если сеть ошиблась, как в случае совы (она не верно указала вид, однако это может быть связано с наличием только такого лейбла в датасете), то её ответ очень близок к истине.

Теперь вспомним, что у нас есть свой датасет, который мы разметили. Для примера возьмём реальные картинки из трёх классов нашего приложения iFunny. Напомню их:

• approved — картинки идут в раздел приложения collective;

• not suitable — не попадают в общую ленту, но остаются в ленте пользователя. К этому относятся девушки в купальниках и мужчины в плавках, селфи и всё, что не является мемами и не несет в себе развлекательную функцию;

• risked — сюда относится расизм, порно, расчленёнка и всё, что законадательно запрещено к размещению и может навредить имиджу компании. Такой контент получает бан и перестает быть доступным всем пользователям iFunny.

В процессе обучения придётся часто обращаться к данным и производить с ними математические операции на ЦПУ или ГПУ (в зависимости от железа). В фреймворке pytorch уже есть реализованный класс ImageFolder, открывающий изображение по заданному пути и присваивающий ему класс в соответствии с папкой, в которой находится. Чтобы им воспользоваться, необходимо сгруппировать изображения определённым образом. В нашем случае все тренировочные изображения лежат в train/interim и разбиты по папкам с названием класса, как это показано ниже. Название объекта должно быть уникальным (у нас это ID контента).

interim/
    approved/
        H6f8XI2i8.jpg
        XFkQE1Zi8.jpg
        ...
    not_suitable/
        DCS2iR3i8.jpg
        KmyGT7Yi8.jpg
        ...
    risked/
        KRXZUuci8.jpg
        m6CH7yxh8.jpg
        ...

Создадим словарь с датасетами тренировочной и валидационной выборок:

from torchvision import datasets

datasets = {
    'train': ImageFolder(
        root='train/interim/',
        transform=transform
    ),
    'valid': ImageFolder(
        root='test/interim/',
        transform=transform
    ),
}

datasets['train'].class_to_idx

Как было сказано ранее, класс ImageFolder производит лейблирование — каждой названной папке сопоставляет определённую цифру. По факту выдает порядковый номер отсортированного по алфавиту списка классов:

{'approved': 0, 'not_suitable’: 1, 'risked’: 2}

Это необходимо, поскольку производить математические операции с числами в процессе оптимизации сети проще, чем обращаться к строкам. Но загружать по одному изображению долго, поэтому используем DataLoader, который делает это сразу пачкой (в DS среде она называется батчом):

torchfrom torch.utils.data import DataLoader

batch_size = 100
num_workers = 20
dataloaders = {
    'train': DataLoader(
        datasets['train'], 
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers
    ),
    'valid': DataLoader(
        datasets['valid'], 
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers
    )
}

Передаём в DataLoader наш датасет и указываем размер батча, который говорит о том, сколько картинок загрузить в сеть одновременно. С точки зрения машины все изображения являются матрицами, а все операции внутри сети — матричными. Поэтому ничего не мешает производить их сразу с N изображениями, собрав их вместе.

Также DataLoader может загружать данные параллельно, за счёт чего процесс загрузки происходит в разы быстрее. Количество одновременных процессов на загрузку задаётся параметром num_workers.

Отобразим наш тренировочный датасет: 

import numpy as np

real_batch = next(iter(dataloaders['train']))
plt.figure(figsize=(8,8))
plt.axis("off")
plt.title("Training Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0][:64], padding=2, normalize=True),(1,2,0)));

Примерно так выглядит контент, который загружают пользователи в iFunny, с единственным отличием, предполагающим сбалансированность классов в нашем датасете, в отличие от реальности, где risked контент составляет меньше 10%:

Перестройка модели под наши цели

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

В следующей строке повторим загрузку модели с заданной архитектурой VGG-11. Флаг pretrained в положении True позволяет загрузить предобученные на ImageNet веса, любезно предоставленные pytorch. Модель переводим в память ГПУ методом to, аргументом которого является название необходимого девайса, так как все дальнейшие вычисления будут производиться на ней. 

device = 'cuda'

model = models.vgg11(pretrained=True, progress=False).to(device)

Затем в цикле присваиваем атрибуту requires_grad всех слоев сети значение False, чтобы в процессе обучения они не изменялись: 

for param in model.parameters():
    param.requires_grad = False

После чего меняем слой, отвечающий за классификацию:

model.classifier[6] = torch.nn.Linear(4096, 3).to(device)

Важный момент: мы не переобучаем все предыдущие слои, так как считаем, что они уже научились выделять общие признаки, в отличие от последнего слоя, который только что поменяли и тем самым переключили его атрибут requires_grad в положение True. Этот слой необходимо обучить, так как теперь он имеет случайные веса, а значит не может корректно отличать что-либо. Также нужно задать оптимизатор, который будет подбирать наиболее подходящие веса и фактически обучать сеть.

Для этого возьмём один из методов градиентного спуска Adam. Ещё необходимо указать критерий, по которому будет идти оптимизация — для этого используем кросс-энтропию.

params_to_update = model.parameters()

print("Params to learn:")
params_to_update = []
for name,param in model.named_parameters():
    if param.requires_grad == True:
        params_to_update.append(param)

optimizer = torch.optim.Adam(params_to_update, lr=0.0001)
criterion = torch.nn.CrossEntropyLoss()

В цикле выше собираем слои, у которых флаг requires_grad в положении True. До этого мы их переводили в положении False, поэтому в оптимизатор передаётся информация только о последнем слое.

Обучение модели

Для обучения воспользуемся функцией train_model из официального туториала pytorch:

import time
import copy

from tqdm import tqdm

def train_model(model, dataloaders, criterion, optimizer, num_epochs=10):
    since = time.time()

    val_acc_history = []

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'valid']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in tqdm(dataloaders[phase]):
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    loss = criterion(outputs, labels)

                    _, preds = torch.max(outputs, 1)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'valid' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
            if phase == 'valid':
                val_acc_history.append(epoch_acc)

        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model, val_acc_history

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

Нужно обязательно обнулять градиент командой optimizer.zero_grad() перед каждым запуском обратного распространения ошибки внутри сети:

loss.backward()
optimizer.step()

Первая строка запускает операцию обратного распространения ошибки внутри сети из переменной потери (loss). Вторая — выполняет градиентный шаг на основе вычисленных градиентов.

Ответы нашей модели считаются в следующей строке:

outputs = model(inputs)

А ошибку, которую используем в дальнейшем для расчёта градиента, получаем одной строчкой:

loss = criterion(outputs, labels)

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

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

if phase == 'train':
model.train()  # Set model to training mode
else:
model.eval()   # Set model to evaluate mode

Для иллюстрации процесса обучения и его результатов, мы провели обучение модели на 10 эпохах. Эпохой в DS среде называется полный проход по всем примерам выборки (обычно сеть не способна за один раз усвоить все правила, поэтому количество эпох почти всегда больше 1). Ниже приведены выводы функции обучения:

model, val_acc_history = train_model(model, dataloaders, criterion, optimizer, num_epochs=10)

0%|          | 0/245 [00:00<?, ?it/s]
Epoch 0/9
----------

100%|██████████| 245/245 [00:51<00:00,  4.76it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.8573 Acc: 0.6141

100%|██████████| 30/30 [00:08<00:00,  3.51it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8701 Acc: 0.5867

Epoch 1/9
----------

100%|██████████| 245/245 [00:51<00:00,  4.73it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7892 Acc: 0.6526

100%|██████████| 30/30 [00:08<00:00,  3.57it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8458 Acc: 0.6012

Epoch 2/9
----------

100%|██████████| 245/245 [00:51<00:00,  4.75it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7716 Acc: 0.6601

100%|██████████| 30/30 [00:08<00:00,  3.50it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8380 Acc: 0.6049

Epoch 3/9
----------

100%|██████████| 245/245 [00:52<00:00,  4.70it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7551 Acc: 0.6658

100%|██████████| 30/30 [00:08<00:00,  3.59it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8374 Acc: 0.6012

Epoch 4/9
----------

100%|██████████| 245/245 [00:52<00:00,  4.71it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7464 Acc: 0.6703

100%|██████████| 30/30 [00:08<00:00,  3.47it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8166 Acc: 0.6157

Epoch 5/9
----------

100%|██████████| 245/245 [00:52<00:00,  4.71it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7423 Acc: 0.6731

100%|██████████| 30/30 [00:08<00:00,  3.57it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8155 Acc: 0.6174

Epoch 6/9
----------

100%|██████████| 245/245 [00:52<00:00,  4.69it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7379 Acc: 0.6764

100%|██████████| 30/30 [00:08<00:00,  3.54it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8117 Acc: 0.6221

Epoch 7/9
----------

100%|██████████| 245/245 [00:52<00:00,  4.69it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7329 Acc: 0.6780

100%|██████████| 30/30 [00:08<00:00,  3.55it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8113 Acc: 0.6201

Epoch 8/9
----------

100%|██████████| 245/245 [00:51<00:00,  4.77it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7314 Acc: 0.6802

100%|██████████| 30/30 [00:08<00:00,  3.50it/s]
  0%|          | 0/245 [00:00<?, ?it/s]
valid Loss: 0.8106 Acc: 0.6221

Epoch 9/9
----------

100%|██████████| 245/245 [00:51<00:00,  4.73it/s]
  0%|          | 0/30 [00:00<?, ?it/s]
train Loss: 0.7243 Acc: 0.6787

100%|██████████| 30/30 [00:08<00:00,  3.65it/s]
valid Loss: 0.8184 Acc: 0.6123

Training complete in 10m 4s
Best val Acc: 0.622095

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

Результаты обучения

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

Сеть хорошо научилась разделять approved и not suitable контент, а вот с изображениями класса risked не всё так гладко. Но напомню, что использовалась очень простенькая сеть без каких-либо дополнительных методов улучшения, поэтому её ещё есть куда развивать.  

Паттерны новой сети

Отразим основные паттерны нашей сети с помощью вектора Шепли.

На первой картинке с котом видно, что сеть при выборе approved класса в основном смотрела на мордочку и уши. У Гитлера определённые паттерны при выборе класса risked выделить сложнее — сеть немного отреагировала на нос и усы, но скорее всего решение принималось по цветовой гамме, так как нацистские фото почти всегда чёрно-белые, и она цепляется за этот признак. 

Самое интересное можно наблюдать на третьей картинке с девушкой. Грудь была важным параметром для выбора обоих классов (и risked, и not suitable), но оказала негативное влияние на класс approved, отклонив его. В пользу not suitable сеть склонил купальник, контур которого у этого класса выделен красным, а у risked — синим. В данном случае все логично, поскольку без купальника данный контент стал бы эротикой и считался недопустимым. Также наличие лица в кадре отрицательно сказалось при проверке сетью принадлежности объекта к risked классу — мы не запрещаем пользователям загружать селфи, но и не продвигаем его через общую ленту, так как такой контент больше подходит для Instagram.

Вместо заключения

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

С помощью метода из первой статьи можно получить первичную разметку, которую придётся перепроверить вручную для лучшего качества, но это будет проще, чем размечать все объекты с нуля. На основе последовательности действий из статьи можно дообучить не только представленную сеть, но и что-то более сложное. 

На данный момент в открытом доступе лежит большое количество готовых реализаций архитектур от Google, Facebook, а также других компаний и университетов с предобученными весами, показывающими отличные результаты, а порой даже относящихся к State of the Art (SOTA). Для поиска подходящей архитектуры есть удобный сайт Papers with Code, где в открытом доступе на GitHub есть готовые реализации.

Напоследок — мем, который был одобрен нашей обученной сетью с вероятностью 85%.

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

approved — 0.8528
not suitable — 0.0996
risked — 0.0475