Время прочтения: 8 мин.

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

Нейронные сети как один из наиболее популярных алгоритмов машинного обучения успешно применяются во многих отраслях деятельности человека. Параллельно с этим совершенствуются и сами технологии для создания и обучения нейросетей. Так в 2017 году состоялся релиз OpenSource-фреймворка PyTorch, который на сегодняшний день является одним из самых популярных в области Deep Learning. С помощью него достаточно легко реализуются простые модели, однако, сложности могут начаться при добавлении дополнительных функций, таких как обучение на GPU (графический процессор) или TPU (тензорный процессор). Здесь на помощь разработчику приходит представленный в марте 2019 года PyTorch Lightning –  библиотека Python с открытым исходным кодом, которая предоставляет высокоуровневый интерфейс для PyTorch.

Один из пунктов Zen of python гласит Simple is better than complex (Простое лучше, чем сложное, англ.). Именно этим принципом руководствовался Уиллиам Фалкон при разработке PyTorch Lightning.

В рамках статьи на примерах будет рассмотрено применение PyTorch и PyTorch Lightning для создания классификатора. В качестве датасета использован известный набор данных MNIST, состоящий из изображений рукописных цифр, представляющих собой одноканальные картинки размером 28х28 пикселей. Задача нейросети научиться определять написанную цифру.

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

Итак, перейдем к примерам. Ниже приведен работоспособный код, который можно протестировать самостоятельно в Jupyter Notebook или Google Colab.

Для начала работы импортируем нужные библиотеки.

import os

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms

import pytorch_lightning as pl
from pytorch_lightning import Trainer

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

EPOCHS = 10
torch.random.manual_seed(42)
pl.seed_everything(42)
Global seed set to 42
42

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

transforms_set=transforms.Compose([transforms.ToTensor(),
                                   transforms.Normalize((0.1307,),
                                                       (0.3081,))])

Загрузим датасет из модуля torchvision.datasets. В параметрах следует указать значение для параметра train. При train=True будет загружен тренировочный датасет, при train=False тестовый (тренировочный датасет можно при необходимости дополнительно разделить на train и validation). Здесь же укажем набор трансформаций.

# Загрузка тренировочного датасета.
train_data = MNIST(root="data",
                   train=True,
                   download=True,
                   transform=transforms_set)

# Загрузка тестового датасета
test_data = MNIST(root="data",
                  train=False,
                  download=True,
                  transform=transforms_set)

Для проверки корректности данных можно проверить результат.

labels_dict = {
    0: "0",
    1: "1",
    2: "2",
    3: "3",
    4: "4",
    5: "5",
    6: "6",
    7: "7",
    8: "8",
    9: "9",
}

fig = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_index = torch.randint(len(train_data), size=(1,)).item()
    img, label = train_data[sample_index]
    fig.add_subplot(rows, cols, i)
    plt.title(labels_dict[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

Затем данные необходимо передать в загрузчик данных DataLoader – итератор из API PyTorch, позволяющий, в том числе загружать данные порциями или батчами. В параметре batch_size следует указать желаемый размер батча, выраженный в количестве изображений.

train_data_loader = DataLoader(train_data, batch_size=16)
test_data_loader = DataLoader(test_data, batch_size=16)

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

Для примера выбрана тривиальная полносвязная нейросеть состоящая из трех слоев. С функцией активации ReLU после входного и скрытого слоев, с функцией Softmax после выходного слоя. Не смотря на простоту, выбранной архитектуры будет достаточно для получения результатов.

Объявляем класс модели и формируем ее архитектуру.

class Pytorch_MNIST_Classifier(nn.Module):

  def __init__(self):
    super(Pytorch_MNIST_Classifier, self).__init__()


    # картинки датасета MNIST  c размером 
     (1, 28, 28) (количество каналов, ширина, высота) 
    self.layer_1 = torch.nn.Linear(28 * 28, 64)
    self.layer_2 = torch.nn.Linear(64, 32)
    self.layer_3 = torch.nn.Linear(32, 10)
 
  def forward(self, x):
    batch_size, channels, width, height = x.size()

    # (b, 1, 28, 28) -> (b, 1*28*28)
    x = x.view(batch_size, -1)

    # слой 1
    x = self.layer_1(x)
    x = torch.relu(x)

    # слой 2
    x = self.layer_2(x)
    x = torch.relu(x)

    # слой 3
    x = self.layer_3(x)

    # вероятность принадлежности к классу
    x = torch.log_softmax(x, dim=1)

    return x

Вызовем объект класса для проверки.

Pytorch_MNIST_Classifier()
Pytorch_MNIST_Classifier(
  (layer_1): Linear(in_features=784, out_features=64, bias=True)
  (layer_2): Linear(in_features=64, out_features=32, bias=True)
  (layer_3): Linear(in_features=32, out_features=10, bias=True)
)

Объявим оптимизатор и функцию потерь.

Pytorch_MNIST_model = Pytorch_MNIST_Classifier()
# Оптимизатор и функция потерь
optimizer = torch.optim.Adam(Pytorch_MNIST_model.parameters(), lr=1e-3)
loss_func = nn.CrossEntropyLoss()

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

# Запускается тренировочный цикл
for epoch in range(EPOCHS):
  size = len(train_data_loader.dataset)
# Загрузчик данных передает данные и их классы
  for batch, (X, y) in enumerate(train_data_loader):

        # Нейросеть обрабатывает данные и возвращает предсказания
        pred = Pytorch_MNIST_model(X)
        # Вычисляется функция потерь
        loss = loss_func(pred, y)

        # Обратным распространением обновляются веса в нейронах
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Для наглядности печатаются значения функции потерь
        if batch % 1000 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

…
loss: 0.009693  [    0/60000]
loss: 0.049324  [16000/60000]
loss: 0.004751  [32000/60000]
loss: 0.002540  [48000/60000]

Нейросеть обучилась на 60 000 изображений. Теперь с ее помощью можно получить предсказания для тестовой выборки и вычислить значение метрики качества модели, для этого применим Accuracy  – долю правильно угаданных классов объектов среди всех предсказаний.

Для тестирования модели также создается цикл, но предварительно модель необходимо «переключить» в режим оценки и отключить вычисление градиентов. Затем, как и в предыдущем цикле, загрузчик передает на вход модели батчи данных, а на выходе каждому элементу присваивается метка, затем вычисляется точность работы алгоритма.

 %%time
size = len(test_data_loader.dataset)
batches = len(test_data_loader)
test_loss, correct = 0, 0

# Модель переключается в режим оценки
Pytorch_MNIST_model.eval()

# Отключение вычисления градиентов
with torch.no_grad():
    # Запускается цикл тестирования
    for X, y in test_data_loader:
        # Нейросеть обрабатывает данные и возвращает предсказания
        pred = Pytorch_MNIST_model(X)
        # Складываются корректные предсказания
        test_loss += loss_func(pred, y).item()
        correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= batches
correct /= size
# Вычисляется метрика качества
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%,
        Avg loss: {test_loss:>8f} \n")

Test Error:
 Accuracy: 96.5%, Avg loss: 0.169451

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

class Pytorch_Lightning_MNIST_Classifier(pl.LightningModule):
  
  def __init__(self):
    super().__init__()
    # Задается архитектура нейросети
    self.layers = nn.Sequential(
      nn.Linear(28 * 28 * 1, 64),
      nn.ReLU(),
      nn.Linear(64, 32),
      nn.ReLU(),
      nn.Linear(32, 10)
    )
    # Объявляется функция потерь
    self.loss_func = nn.CrossEntropyLoss()
    
  def forward(self, x):
    return self.layers(x)
  
  # Настраиваются параметры обучения
  def training_step(self, batch, batch_idx):
    x, y = batch
    x = x.view(x.size(0), -1)
    pred = self.layers(x)
    loss = self.loss_func(pred, y)
    self.log('train_loss', loss)
    return loss
  
  # Настраиваются параметры тестирования
  def test_step(self, batch, batch_idx):
    x, y = batch
    x = x.view(x.size(0), -1)
    pred = self.layers(x)
    loss = self.loss_func(pred, y)
    pred = torch.argmax(pred, dim=1)
    accuracy = torch.sum(y == pred).item() / (len(y) * 1.0)
    self.log('test_loss', loss, prog_bar=True)
    self.log('test_acc', torch.tensor(accuracy), prog_bar=True)
    output = dict({
        'test_loss': loss,
        'test_acc': torch.tensor(accuracy),
    })
    return output

  # Конфигурируется оптимизатор
  def configure_optimizers(self):
    optimizer = torch.optim.Adam(self.parameters(), lr=1e-4)
    return optimizer

Для обучения и тестирования модели в PyTorch Lightning предусмотрена специальная функция Trainer(), имеющая широкий набор настроек, к примеру, можно указать вручную количество графических процессоров, если мы имеем дело с распределенными вычислениями. Перед началом обучения Trainer() самостоятельно проверит наличие GPU или TPU, и переключит вычисления на доступный, в PyTorch это надо прописывать вручную.

В Trainer() добавлен ряд привычных из классического ML методов, таких как fit, test, predict, доступных «из коробки».

# Инициализация модели и функции Trainer
Pytorch_lightning_MNIST_model = Pytorch_Lightning_MNIST_Classifier()
trainer = pl.Trainer(gpus=None,
                     max_epochs=EPOCHS)

# Обучение модели
trainer.fit(Pytorch_lightning_MNIST_model, DataLoader(train_data))

# Тестирование модели
trainer.test(Pytorch_lightning_MNIST_model, DataLoader(test_data))

GPU available: False, used: False
TPU available: False, using: 0 TPU cores

  | Name      | Type             | Params
-----------------------------------------------
0 | layers    | Sequential       | 52.6 K
1 | loss_func | CrossEntropyLoss | 0     
-----------------------------------------------
52.6 K    Trainable params
0         Non-trainable params
52.6 K    Total params
0.211     Total estimated model params size (MB)
Epoch 9: 100%
60000/60000 [03:33<00:00, 280.51it/s, loss=0.000677, v_num=0]

Testing: 100%
10000/10000 [00:11<00:00, 847.06it/s]

--------------------------------------------------------------------------
DATALOADER:0 TEST RESULTS
{'test_acc': 0.9688000082969666, 'test_loss': 0.13828878104686737}
--------------------------------------------------------------------------

Количество и разнообразие задач, решаемых путем применения моделей глубокого обучения растет ежедневно, вместе с этим, растет потребность в мощных и доступных инструментах для их решения, таких как PyTorch. Но индустрия не стоит на месте и сегодня PyTorch Lightning не уступает PyTorch в спектре возможностей, имея при этом ряд преимуществ:

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

Все вышеперечисленное позволяет использовать фреймворк PyTorch Lightning наравне с другими библиотеками для решения прикладных и исследовательских задач.