Практическая работа №4. Автоматическая оптимизация нейронной сети с помощью Apache TVM (дополнительные материалы)¶
Данный скрипт содержит процедуру обучения сетей, соответствующих логистической регрессии и многослойной полносвязной сети, для решения задачи классификации рукописных цифр на наборе данных MNIST. Модель обучается с использованием фреймворка PyTorch. По завершении обучения модель конвертируется и сохраняется в формате тензорного компилятора Apache TVM для последующей оптимизации слоев.
В скрипте выполняется стандартная последовательность действий:
- Загрузка набора данных MNIST и проверка корректности загруженных изображений и меток.
Предполагает использование возможностей пакета
torchvision.datasets
. - Формирование архитектуры модели логистической регрессии. Предполагает реализацию класса
LogisticRegressionModel
, который является наследником классаtorch.nn.Module
. - Реализация функции обучения нейронной сети средствами PyTorch.
- Тестирование модели с использованием средств библиотеки PyTorch.
- Формирование архитектуры многослойной полносвязной сети, ее обучение и тестирование по аналогии с логистической регрессией.
- Формирование архитектуры сверточной нейронной сети, ее обучение и тестирование по аналогии с логистической регрессией.
- Конвертация и сохранение модели в формат тензорного компилятора TVM.
1. Импорт пакетов и загрузка данных¶
In [1]:
import torch
import torchmetrics
import torchvision
import torchvision.datasets
import torchvision.transforms
import torch.utils.data
import torch.nn
import torch.nn.functional as F
import os
from matplotlib import pyplot as plot
from time import time
In [2]:
train_dataset = torchvision.datasets.MNIST(
root = '.', train = True, download = True,
transform = torchvision.transforms.ToTensor()
)
test_dataset = torchvision.datasets.MNIST(
root = '.', train = False, download = True,
transform = torchvision.transforms.ToTensor()
)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
In [3]:
def show_images(images, title):
num_showed_imgs_x = 5
num_showed_imgs_y = 5
figsize = (2, 2)
fig, axes = plot.subplots(num_showed_imgs_y, num_showed_imgs_x, figsize = figsize)
fig.suptitle(title)
plot.setp(plot.gcf().get_axes(), xticks = [], yticks = [])
for i, ax in enumerate(axes.flat):
img = images[i][0].numpy().transpose(1, 2, 0).squeeze(axis = 2)
ax.imshow(img, cmap = 'gray')
print('Number of train samples: {}'.format(len(train_dataset)))
show_images(train_dataset, 'Train samples')
print('Number of test samples: {}'.format(len(test_dataset)))
show_images(test_dataset, 'Test samples')
train_data_loader = torch.utils.data.DataLoader(
train_dataset, batch_size = 64, shuffle = True
)
test_data_loader = torch.utils.data.DataLoader(
test_dataset, batch_size = 64, shuffle = False
)
Number of train samples: 60000 Number of test samples: 10000
2. Формирование архитектуры модели логистической регресии¶
In [4]:
image_resolution = 28 * 28
num_classes = 10
class LogisticRegressionModel(torch.nn.Module):
def __init__(self, input_dim, output_dim):
super(LogisticRegressionModel, self).__init__()
self.linear = torch.nn.Linear(input_dim, output_dim)
def forward(self, x):
out = self.linear(x)
return out
logreg_nn = LogisticRegressionModel(image_resolution, num_classes)
3. Обучение модели¶
In [5]:
def fcnn_transform_data(images):
return images.view(-1, image_resolution)
def train_model(model, train_data_loader, transform_data, num_epochs=10):
model.to(device)
loss_function = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=10)
metric.to(device)
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_data_loader):
images = transform_data(images).requires_grad_().to(device)
labels = labels.to(device)
outputs = model(images)
loss = loss_function(outputs, labels)
accuracy = metric(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
accuracy = metric.compute()
print(f'Epoch[{epoch}]: train accuracy = {accuracy:.3f}')
return
In [6]:
train_model(
logreg_nn, train_data_loader, fcnn_transform_data,
num_epochs=10
)
Epoch[0]: train accuracy = 0.870 Epoch[1]: train accuracy = 0.891 Epoch[2]: train accuracy = 0.900 Epoch[3]: train accuracy = 0.905 Epoch[4]: train accuracy = 0.909 Epoch[5]: train accuracy = 0.912 Epoch[6]: train accuracy = 0.914 Epoch[7]: train accuracy = 0.916 Epoch[8]: train accuracy = 0.917 Epoch[9]: train accuracy = 0.918
4. Тестирование модели¶
In [7]:
def get_accuracy(data_loader, transform_data, model):
model.to(device)
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=10)
metric.to(device)
for i, (images, labels) in enumerate(data_loader):
images = transform_data(images).requires_grad_().to(device)
labels = labels.to(device)
outputs = model(images)
accuracy = metric(outputs, labels)
return metric.compute()
In [8]:
logreg_accuracy = get_accuracy(test_data_loader, fcnn_transform_data, logreg_nn)
print(f'Test accuracy: {logreg_accuracy:.3f}')
Test accuracy: 0.926
5. Формирование архитектуры многослойной полносвязной сети, обучение и тестирование¶
In [9]:
class FCNNModel(torch.nn.Module):
def __init__(self, input_dim, output_dim, hidden_dim):
super(FCNNModel, self).__init__()
self.linear1 = torch.nn.Linear(input_dim, hidden_dim)
self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
self.linear3 = torch.nn.Linear(hidden_dim, hidden_dim)
self.linear4 = torch.nn.Linear(hidden_dim, output_dim)
def forward(self, x):
x = F.relu(self.linear1(x))
x = F.relu(self.linear2(x))
x = F.relu(self.linear3(x))
x = self.linear4(x)
return x
fcnn = FCNNModel(image_resolution, num_classes, hidden_dim=300)
In [10]:
train_model(
fcnn, train_data_loader, fcnn_transform_data,
num_epochs=10
)
Epoch[0]: train accuracy = 0.920 Epoch[1]: train accuracy = 0.944 Epoch[2]: train accuracy = 0.956 Epoch[3]: train accuracy = 0.963 Epoch[4]: train accuracy = 0.968 Epoch[5]: train accuracy = 0.971 Epoch[6]: train accuracy = 0.974 Epoch[7]: train accuracy = 0.976 Epoch[8]: train accuracy = 0.978 Epoch[9]: train accuracy = 0.980
In [11]:
fcnn_accuracy = get_accuracy(test_data_loader, fcnn_transform_data, fcnn)
print(f'Test accuracy: {fcnn_accuracy:.3f}')
Test accuracy: 0.980
6. Формирование архитектуры сверточной сети, обучение и тестирование¶
In [12]:
def cnn_transform_data(images):
return images
class CNNModel(torch.nn.Module):
def __init__(self, input_channels, conv_filters, num_classes):
super(CNNModel, self).__init__()
self.conv = torch.nn.Conv2d(input_channels, conv_filters, 3, padding=1)
self.pool = torch.nn.MaxPool2d(2)
self.linear = torch.nn.Linear(14*14*conv_filters, num_classes)
def forward(self, x):
x = self.pool(F.relu(self.conv(x)))
x = x.view(x.shape[0], -1)
x = self.linear(x)
return x
cnn = CNNModel(input_channels=1, conv_filters=64, num_classes=num_classes)
In [13]:
train_model(
cnn, train_data_loader, cnn_transform_data,
num_epochs=10
)
Epoch[0]: train accuracy = 0.943 Epoch[1]: train accuracy = 0.961 Epoch[2]: train accuracy = 0.969 Epoch[3]: train accuracy = 0.973 Epoch[4]: train accuracy = 0.977 Epoch[5]: train accuracy = 0.979 Epoch[6]: train accuracy = 0.981 Epoch[7]: train accuracy = 0.982 Epoch[8]: train accuracy = 0.984 Epoch[9]: train accuracy = 0.985
In [14]:
cnn_accuracy = get_accuracy(test_data_loader, cnn_transform_data, cnn)
print(f'Test accuracy: {cnn_accuracy:.3f}')
Test accuracy: 0.985
7. Конвертация и сохрание моделей в формат Apache TVM¶
Ниже реализована функция конвертации. Наряду с этим, выполняется сохранение изображений и меток в формате NumPy для упрощения тестирования корректности работы моделей на RISC-V.
In [15]:
import tvm
from tvm import relay
import numpy as np
os.makedirs('model/', exist_ok=True)
def convert_to_tvm_and_save(model, name, input_shape):
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model.to('cpu'), input_data).eval()
shape_list = [('input0', input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
with open(f'model/{name}.params', "wb") as fo:
fo.write(relay.save_param_dict(params))
with open(f'model/{name}.json', "w") as fo:
fo.write(tvm.ir.save_json(mod))
convert_to_tvm_and_save(logreg_nn, 'logreg', input_shape=(1, image_resolution))
convert_to_tvm_and_save(fcnn, 'fcnn', input_shape=(1, image_resolution))
convert_to_tvm_and_save(cnn, 'cnn', input_shape=(1, 1, 28, 28))
result = {
'logreg': logreg_accuracy.detach().cpu().numpy(),
'fcnn': fcnn_accuracy.detach().cpu().numpy(),
'cnn': cnn_accuracy.detach().cpu().numpy(),
}
np.save('model/metric.npy', result)
In [16]:
os.makedirs('data/', exist_ok=True)
test_data_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=len(test_dataset), shuffle=False
)
images, labels = next(iter(test_data_loader))
images_np = images.numpy()
labels_np = labels.numpy()
np.save('data/test_images.npy', images_np)
np.save('data/test_labels.npy', labels_np)