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

В этом посте вы узнаете:

• Как сохранить модель PyTorch в формате обмена данными.

• Как использовать Netron для создания графического представления.

Давайте начнем.

Почему графическое представление модели PyTorch затруднительно

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

Это проблема: имея модель вы никогда не узнаете, как она работает, пока не проследите за входным тензором и не соберете трассировку до получения выходного тензора. Поэтому преобразование модели PyTorch в картинку не является тривиальным.

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

Как использовать Netron для графического представления модели

Когда вы сохраняете модель PyTorch, вы сохраняете ее состояния. Вы можете получить состояния модели с помощью функции model.state_dict().

Хотя тензоры весов имеют имена, которые помогут вам восстановить их в модели, у вас нет подсказок о том, как весовые коэффициенты связаны друг с другом. 

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

Действительно, если вы хотите узнать алгоритм, лежащий в основе модели PyTorch, это тоже возможный вариант. 

Существует всего несколько инструментов для создания графического представления для модели PyTorch. 

Ниже вы узнаете об инструменте Netron. Это «средство визуализации моделей глубокого обучения». Это программное обеспечение, которое можно установить и запустить на macOS, Linux и Windows. Вы можете посетить страницу ниже и загрузить программное обеспечение для своей платформы:

• https://github.com/lutzroeder/netron/releases

Существует также онлайн-сервис, в котором вы можете визуализировать свою модель, загрузив файл модели.

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

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

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

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

data = load_iris()
X = data['data']
y = data['target']
X = torch.tensor(X, dtype=torch.float32)
y = torch.tensor(y, dtype=torch.long)

# split
X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.7, shuffle=True)

class IrisModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(4, 8)
        self.act = nn.ReLU()
        self.output = nn.Linear(8, 3)

    def forward(self, x):
        x = self.act(self.hidden(x))
        x = self.output(x)
        return x

# loss metric and optimizer
model = IrisModel()
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# prepare model and training parameters
n_epochs = 100
batch_size = 10
batch_start = torch.arange(0, len(X_train), batch_size)

# training loop
for epoch in range(n_epochs):
    for start in batch_start:
        # take a batch
        X_batch = X_train[start:start+batch_size]
        y_batch = y_train[start:start+batch_size]
        # forward pass
        y_pred = model(X_batch)
        loss = loss_fn(y_pred, y_batch)
        # backward pass
        optimizer.zero_grad()
        loss.backward()
        # update weights
        optimizer.step()

# validating model
y_pred = model(X_test)
acc = (torch.argmax(y_pred, 1) == y_test).float().mean()
acc = float(acc)*100
print("Model accuracy: %.2f%%" % acc)

Выполнение вышеуказанного кода дает, например, следующее:

Model accuracy: 97.78%

Итак, вы знаете, что model это модель PyTorch, которая может принимать тензор и возвращать тензор. Вы можете преобразовать эту модель в формат ONNX с помощью функции torch.onnx.export():

torch.onnx.export(model, X_test, 'iris.onnx', input_names=["features"], output_names=["logits"])

Выполнение этой функции создаст файл iris.onnx в локальном каталоге. Вам необходимо предоставить пример тензора, который работает с моделью в качестве входных данных (X_test в приведенном выше примере). 

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

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

Обычно функцию export() следует запускать после цикла обучения. Это связано с тем, что созданная модель ONNX содержит полную модель, которую можно запустить без библиотеки PyTorch. Вы хотите сохранить оптимизированный вес в ней. Однако для целей визуализации модели в Netron качество модели не имеет значения. Вы можете вызвать функцию export(), как только модель PyTorch будет создана.

После запуска Netron вы можете открыть сохраненный файл ONNX. В данном примере вы должны увидеть следующий экран:

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

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

На приведенном выше экране вы видите, что слой nn.Linear() становится «Gemm», что означает «операция умножения общей матрицы». Вы даже можете исследовать веса на слое с помощью Netron, поместив на них курсор и щелкнув кнопкой мыши.

Если вы хотите сохранить копию этой визуализации, вы можете экспортировать ее в формат PNG в Netron.

Дополнительные материалы

Netron — это проект с открытым исходным кодом, исходный код которого вы можете найти на Github:

• https://github.com/lutzroeder/netron

Онлайн версия Netron доступна по ссылке:

• https://netron.app/

Еще одна библиотека для визуализации — torchviz, но в отличие от примера, который вы видели выше, она трассирует модель с помощью обратного прохода:

• https://github.com/szagoruyko/pytorchviz

Резюме

В этом посте вы узнали, как визуализировать модель. В частности, вы узнали:

• Почему визуализация модели PyTorch является сложной задачей.

• Как преобразовать модель PyTorch в формат ONNX.

• Как использовать Netron для визуализации модели ONNX.