Всем привет! Меня зовут Вадим, я Data Scientist в компании Raft, и сегодня мы продолжим погружаться в Mojo. Эта статья является продолжением первой части обзора данного языка программирования. В ней я подробно рассмотрел его преимущества, примеры использования, а также провел сравнение с Python. Если вы ещё не читали прошлую часть, то рекомендую сделать это! 

В этой статье мы погрузимся в практическое применение: обучим линейную регрессию и простую сверточную нейронную сеть, а затем проведём сравнение производительности на этих языках, чтобы понять: действительно ли Mojo превосходит Python в 30 раз?

В качестве задач я выбрал два стандартных соревнования по машинному обучению: предсказание стоимости жилья и классификацию рукописных цифр MNIST. Для проведения экспериментов на Python будет использоваться фреймворк машинного обучения - PyTorch, а на Mojo - Basalt, который описывался в первой части.

Примечание: если вас интересует только производительность Mojo в сравнении с Python, вы сможете найти эту информацию в последней секции данной статьи.

Итак, давайте начнём!

Немного о датасетах

MNIST (Modified National Institute of Standards and Technology) представляет из себя датасет для задачи распознавания рукописных цифр от 0 до 9. Данные состоят из 70 тысяч картинок с разрешением 28х28, каждая из которых имеет черный фон, на котором изображена цифра белого цвета. Задача состоит в том, чтобы распознать цифру, которая изображена на картинке. 

Housing Prices Dataset представляет из себя набор данных для предсказания стоимости жилья на основе некоторых признаков, например, площадь участка, тип жилья, наличие гаража, количество комнат и так далее.

Погружаемся в код

Эксперимент на MNIST

Для решения задачи классификации рукописных цифр напишем простую CNN (convolutional neural network), которая будет состоять из 2-х частей:

• построение карты признаков (feature map), реализованной через 2 слоя сверток

• классификатор, состоящий из 3-х полносвязных слоев.

Более подробно архитектура представлена в таблице ниже.

Гиперпараметры обучения:

• num_epochs = 20

• batch_size = 8

• learning_rate = 2e-3

• оптимизатор Adam

• функция потерь CrossEntropyLoss.

Реализация архитектуры сети на Python и Mojo немного отличается. В первом случае, используя PyTorch, мы могли бы определить архитектуру как последовательность блоков. 

class CNN(nn.Module):

    def __init__(self):

        super(CNN, self).__init__()

        self.block1 = nn.Sequential(

            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, padding=2),

            nn.ReLU(),

            nn.MaxPool2d(kernel_size=2),

        )

        self.block2 = nn.Sequential(

            nn.Conv2d(in_channels=16,out_channels= 32, kernel_size=5, padding=2),

            nn.ReLU(),

            nn.MaxPool2d(kernel_size=2),

        )

        self.fc1 = nn.Linear(in_features=32 * 7 * 7, out_features=120)

        self.fc2 = nn.Linear(in_features=120, out_features=84)

        self.out = nn.Linear(in_features=84, out_features=10)

    def forward(self, x):

        x = self.block1(x)

        x = self.block2(x)

        x = x.view(x.size(0), -1)

        x = nn.ReLU()(self.fc1(x))

        x = nn.ReLU()(self.fc2(x))

        return self.out(x)

В случае с Mojo необходимо определить структуру Graph, которая реализует так называемый граф вычислений, применяемый для вычислений в предсказании (feed forward) и обратного распространения ошибки (backpropagation).

fn create_CNN(batch_size: Int) -> Graph:

   # инициализируем граф и наш вход

    var g = Graph()

    var x = g.input(TensorShape(batch_size, 1, 28, 28))

   # инициализируем и применяем сверточные слои

    var conv1 = nn.Conv2d(g, x, out_channels=16, kernel_size=5, padding=2)

    var act_conv1 = nn.ReLU(g, conv1)

    var max_pool1 = nn.MaxPool2d(g, act_conv1, kernel_size=2)

    var conv2 = nn.Conv2d(g, max_pool1, out_channels=32, kernel_size=5, padding=2)

    var act_conv2 = nn.ReLU(g, conv2)

    var max_pool2 = nn.MaxPool2d(g, act_conv2, kernel_size=2)

    # переводим выходной тензор в вектор

    var x_reshape = g.op(

        OP.RESHAPE,

        max_pool2,

        attributes=AttributeVector(

            Attribute(

                "shape",

                TensorShape(max_pool2.shape[0], max_pool2.shape[1] * max_pool2.shape[2] * max_pool2.shape[3]),

            )

        ),

    )

    # классифицируем, извлеченные признаки, полносвязной сетью

    var fc1 = nn.Linear(g, x_reshape, n_outputs=120)

    var act_fc1 = nn.ReLU(g, fc1)

    var fc2 = nn.Linear(g, act_fc1, n_outputs=84)

    var act_fc2 = nn.ReLU(g, fc2)

    var out = nn.Linear(g, act_fc2, n_outputs=10)

    g.out(out)

    # считаем потери, используя CrossEntropyLoss

    var y_true = g.input(TensorShape(batch_size, 10))

    var loss = nn.CrossEntropyLoss(g, out, y_true)

    g.loss(loss)

    return g

Инициализация модели вместе с оптимизатором и цикл её обучения на Python достаточно стандартны для PyTorch: прогоняем весь датасет некоторое число эпох по батчам, определяем признаки (images) и метки к ним (labels), далее предсказываем класс, рассчитываем ошибку и обновляем градиенты. 

   # определяем модель, функцию потерь и оптимизатор
   cnn = CNN()

   loss_func = nn.CrossEntropyLoss()

   optimizer = optim.Adam(cnn.parameters(), lr=learning_rate)

   cnn.train()

   for epoch in range(num_epochs):

       for i, (images, labels) in enumerate(loaders["train"]):

           b_x = Variable(images)

           b_y = Variable(labels)

          # предсказываем метку класса
           output = cnn(b_x)
          # считаем ошибку
           loss = loss_func(output, b_y)

           optimizer.zero_grad()

           loss.backward()

           optimizer.step()

На Mojo есть небольшие отличия:

1. необходимо определять функцию для выполнения кода

2. необходимо определить модель и оптимизатор через структуру graph

3. перед подачей в изображений в сеть необходимо произвести one hot encoding меток.

В остальном процесс обучения сети схож со стилем PyTorch, за исключением особенностей синтаксиса языка.

fn main():

    alias graph = create_CNN(batch_size)

    var model = nn.Model[graph]()

    var optim = nn.optim.Adam[graph](Reference(model.parameters), lr=learning_rate)

for epoch in range(num_epochs):

        var num_batches: Int = 0

        var epoch_loss: Float32 = 0.0

        for batch in training_loader:

            var labels_one_hot = Tensor[dtype](batch.labels.dim(0), 10)

            for bb in range(batch.labels.dim(0)):

                labels_one_hot[int((bb * 10 + batch.labels[bb]))] = 1.0

            var loss = model.forward(batch.data, labels_one_hot)

            optim.zero_grad()

            model.backward()

            optim.step()

            epoch_loss += loss[0]

            num_batches += 1

House price prediction

Для решения этой задачи мы применим стандартную линейную регрессию, реализованную через один полносвязный слой.

Гиперпараметры обучения следующие:

• num_epochs = 500

• batch_size = 32

• learning_rate = 0.01

• оптимизатор Adam 

• функция потерь MSELoss.

На Python код, с использованием PyTorch,  будет выглядеть следующим образом.

class LinearRegression(nn.Module):

    def __init__(self, input_dim):

        super(LinearRegression, self).__init__()

        self.linear = nn.Linear(in_features=input_dim, out_features=1)

    def forward(self, x):

        return self.linear(x)

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

fn linear_regression(batch_size: Int, n_inputs: Int, n_outputs: Int) -> Graph:

    var g = Graph()

    var x = g.input(TensorShape(batch_size, n_inputs))

    var y_true = g.input(TensorShape(batch_size, n_outputs))

    var y_pred = nn.Linear(g, x, n_outputs)

    g.out(y_pred)

    var loss = nn.MSELoss(g, y_pred, y_true)

    g.loss(loss)

    return g

Цикл обучения совпадает с тем, что был показан на MNIST за исключением того, что необходимость в ohe hot encoding отпадает, так как метки уже закодированы.

Сравнение производительности

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

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

Заключение

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

А что думаете вы? Пишите в комментариях!

P.S: вы можете подписаться на мой телеграм-канал, в котором я освещаю различные темы из AI, а также на блог компании Raft, в нем собрано множество интересных статей.

Ссылки

1. Репозиторий с кодом экспериментов

2. Соревнование на Kaggle о MNIST

3. Соревнование на Kaggle о House Price Prediction

4. Документация PyTorch

5. Документация Mojo

6. Репозиторий Basalt