Практическая работа №4. Автоматическая оптимизация нейронной сети с помощью Apache TVM (дополнительные материалы)¶

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

В скрипте выполняется стандартная последовательность действий:

  1. Загрузка набора данных MNIST и проверка корректности загруженных изображений и меток. Предполагает использование возможностей пакета torchvision.datasets.
  2. Формирование архитектуры модели логистической регрессии. Предполагает реализацию класса LogisticRegressionModel, который является наследником класса torch.nn.Module.
  3. Реализация функции обучения нейронной сети средствами PyTorch.
  4. Тестирование модели с использованием средств библиотеки PyTorch.
  5. Формирование архитектуры многослойной полносвязной сети, ее обучение и тестирование по аналогии с логистической регрессией.
  6. Формирование архитектуры сверточной нейронной сети, ее обучение и тестирование по аналогии с логистической регрессией.
  7. Конвертация и сохранение модели в формат тензорного компилятора 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
No description has been provided for this image
No description has been provided for this image

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)