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

Но Андрей Карпаты, известный исследователь в области ИИ, считает, что реализация алгоритмов с нуля позволяет понять их суть и детали работы, что сложно осознать, используя только готовые библиотеки. Это помогает развить интуицию для дальнейшего применения и улучшения методов. Андрей посвящает много собственного времени, чтобы объяснять ключевые принципы работы нейросетей в своих блогах и на своём ютуб-канале. Он также не раз подчеркивал, что на его курсе в Cтэнфорде есть задачи по реализации различных алгоритмов, например, обратное распространение.

Я хотел бы посвятить данную статью этой идеи, потому что мне самому особенно интересно копаться в алгоритмах глубокого обучения.

Итак, начнём!

Вы наверняка работали с библиотекой PyTorch. Поэтому вам будет знаком этот фрагмент кода

# Создаем простой набор данных
X = torch.randn(100, 3)  # 100 примеров с 3 признаками
y = torch.randint(0, 2, (100,))  # 100 меток классов (0 или 1)

# Определим простую нейронную сеть
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(3, 5)  # Первый слой: 3 входа, 5 выходов
        self.fc2 = nn.Linear(5, 2)  # Второй слой: 5 входов, 2 выхода (классы)
        self.softmax = nn.Softmax(dim=1)  # Для получения вероятностей классов

    def forward(self, x):
        x = torch.relu(self.fc1(x))  # Применяем активацию ReLU
        x = self.fc2(x)  # Второй слой
        x = self.softmax(x)  # Преобразуем в вероятности
        return x

# Создаем модель
model = SimpleNN()

# Определяем функцию потерь и оптимизатор
criterion = nn.CrossEntropyLoss()  # Кросс-энтропия для многоклассовой классификации
optimizer = optim.SGD(model.parameters(), lr=0.01)  # Стохастический градиентный спуск

# Обучаем модель
num_epochs = 100
for epoch in range(num_epochs):
    # Прямой проход
    outputs = model(X)
    
    # Вычисление потерь
    loss = criterion(outputs, y)
    
    # Обратный проход
    optimizer.zero_grad()  # Обнуляем градиенты
    loss.backward()  # Вычисляем градиенты
    optimizer.step()  # Обновляем параметры модели

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

X = torch.randn(100, 3)  # 100 примеров с 3 признаками
y = torch.randint(0, 2, (100,))  # 100 меток классов (0 или 1)

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

Note: Под интуицией я понимаю некоторую идею. Например, если мы будем увеличивать модели, то они смогут решать более сложные задачи - это интуиция.

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(3, 5)  # Первый слой: 3 входа, 5 выходов
        self.fc2 = nn.Linear(5, 2)  # Второй слой: 5 входов, 2 выхода (классы)
        self.softmax = nn.Softmax(dim=1)  # Для получения вероятностей классов

Что это за волшебный модуль nn? Почему мы берем из него готовые блоки и почему наследуемся от nn.Module?

def forward(self, x):
    x = torch.relu(self.fc1(x))  # Применяем активацию ReLU
    x = self.fc2(x)  # Второй слой
    x = self.softmax(x)  # Преобразуем в вероятности
    return x

И почему для процесса обучения модели хватило описания всего 2 методов?

# Определяем функцию потерь и оптимизатор
criterion = nn.CrossEntropyLoss()  # Кросс-энтропия для многоклассовой классификации
optimizer = optim.SGD(model.parameters(), lr=0.01)  # Стохастический градиентный спуск

Почему нам необходимо определять объект от некоего класса для нашей loss-функции? Почему нам необходимо определять объект от некоего класса для градиентного спуска? Почему мы вдруг должны передавать некую сущность model.parameters() в этот объект?

# Вычисление потерь
loss = criterion(outputs, y)
# Обратный проход
optimizer.zero_grad()  # Обнуляем градиенты
loss.backward()  # Вычисляем градиенты
optimizer.step()  # Обновляем параметры модели

С какой целью нам необходимо обнулять градиенты оптимизатора? Почему вычисление градиентов отделено от обновления параметров модели?

Эти и многие другие вопросы доставляли мне дискомфорт, когда я писал код для обучения своей модели. (Вы наверняка знаете то чувство, когда совершаете какое-то действие, но не совсем понимаете зачем оно нужно или как оно работает). Впоследствии я конечно же ответил на свои возникшие вопросы!

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

Поэтому я поставил себе цель: Написать маленькую библиотеку машинного обучения используя NumPy, которая поможет мне разобраться в некоторых тонкостях большой библиотеки.

Теперь, когда я её собрал и она работает, я бы хотел бы с вами поделиться знаниями, которые я получил во время её создания.

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

Конечный ожидаемый результат моей работы: Написать ОБЁРТКУ под NumPy, так чтобы после замены строчек

import torch
import torch.nn as nn

на

import <наша библиотека>
import <наша библиотека>.nn as nn

У нас бы всё продолжило работать и обучаться. Я буду рассказывать темы в том порядке, в котором я их сам изучал. И так как у меня получилось всё соединить в единое целое, то я хотел бы этот порядок соблюсти.

Я подозреваю что вы уже знакомы с PyTorch, NumPy, Gradient Descent и обучали свою модель.

Module

Это основа нашей библиотеки. Любые слои - готовые и созданные нами будут наследоваться от этого класса. Почему? Потому что это упрощает написание кода, так как всё что остается пользователю - написать составную часть блока (например, 2 линейных слоя или 3 свёрточных слоя) и затем определить как будут изменяться данные при прохождении через этот блок.

Давайте напишем скелет этого класса:

class Module:

    def __init__(self):
        pass

    def forward(self):
        pass

    def __call__(self, x):
        return self.forward(x)

Вспомните, как в pytorch: output = model(input) - красиво выглядит, да? Гораздо лучше чем,output = model.forward(input), это достигается путём определения метода __call__(self[, args…]) - Он позволяет экземплярам пользовательских типов представляться объектами, поддерживающими вызов. Но зачем всё-таки мы определяем forward(self, args)?

Суть в том, чтобы отделить логику обработки входного значения пользователем от логики работы самого класса. Всё как в жизни. Инструкцию пишем в одном месте, а как её обрабатывать в другом (причем часто пишет её не пользователь). В добавок прописывать в методе forward(self, args) что-то в родительском классе нам нет необходимости, эта часть полностью ложится на плечи пользователя.

У нас уже получилось создать минимально необходимую часть нейронной сети. Давайте посмотрим на её работу

class Module:
    def __init__(self):
        pass
    def forward(self):
        pass
    def __call__(self, x):
        return self.forward(x)

class Node(Module):
    def __init__(self):
        pass
    def forward(self, x):
        return x ** 2

model = Node()
x = 5
model(x)
>>> 25

Отлично! Давайте продолжим дополнять и усложнять!

Теперь нам нужно подумать, как будет выглядеть наш объект, который будет нести в себе значения

x = 5.0 # - это может быть просто число
x = [1.0, 2.0, 3.0] # - это может быть список чисел
x = [[1.0, 2.0, 3.0],
     [4.0, 5.0, 6.0]] # - это может быть матрица чисел произвольного размера

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

import numpy as np
x = np.array(5) 
>>> array(5)

x = np.array([1.0, 2.0, 3.0])
>>> array([1., 2., 3.])
 
x = np.array([[1.0, 2.0, 3.0],
                [4.0, 5.0, 6.0]]) 
>>> array([[1., 2., 3.],
           [4., 5., 6.]])

Столпом нейронных сетей является линейный слой или по-другому перцептрон. Давайте попробуем реализовать его на основе имеющихся знаний.

class Linear:
    def __init__(self):
        self.weight = 3

    def __call__(self, x):
        result = self.weight * x
        return result

Итого получаем

model = Linear()
x = 5
model(x)
>>> 15

А теперь, зная что число - это всего лишь частный вид матрицы, расширим наш линейный слой на матрицы и матричные умножения. Также добавим bias, чтобы соблюдать каноничность.

class Linear:
    def __init__(self, input_channels, output_channels, bias):
        self.weight = np.random.uniform(- 0.5, 0.5, size=(input_channels, output_channels))
        if bias:
            self.bias = np.random.uniform(- 0.5, 0.5, size=output_channels)
        else:
            self.bias = np.zeros(output_channels)
            
    def __call__(self, x):
        self.x = x # сохраним входное значение на всякий случай
        result = x @ self.weight + self.bias
        return result

Результат

model = Linear(input_channels=3, output_channels=2, bias=True)
x = np.random.randn(3, 3) # (3, 3) - матрица
model(x) # (3, 3) @ (3, 2) + (1, 2) -> (3, 2)
>>> array([[ 0.09383105,  0.32933136],
       [ 0.24610282, -0.20840119],
       [ 0.14503834,  0.51688682]])

Круто! Мы уже можем собрать полноценную многослойную нейронку!

linear1 = Linear(input_channels=5, output_channels=5, bias=True)
linear2 = Linear(input_channels=5, output_channels=3, bias=True)
x = np.random.randn(5, 5) 
x_1 = linear1(x)
x_2 = linear2(x_1) # (5, 3)
>>> array([[ 0.08126928, -0.57829926, -0.21338573],
        [ 0.01379938, -0.36082656,  0.5538853 ],
        [ 0.01338218, -0.29319368,  0.19510403],
        [-0.00940793, -0.55011709,  0.57888208],
        [ 0.02418422, -0.30528372,  0.27401544]]

Классификация

Ну вот мы получили выход с нейронной сети, что нам теперь с ним сделать? Сравнить с эталонным! Например, нам нужно было получить 5, а получили 3, тогда нам нужно изменить веса модели, чтобы она приблизила свой ответ от 3 к 5. Давайте сразу перейдем от задачи регрессии к задаче классификации, так как она более интересная и показательная.

Я подозреваю, что вы знаете задачу классификации и поэтому просто напомню. Задача классификации - выбрать один верный класс из некоторого множества всех классов. Например, что изображено на картинке - кот, собака, самолет или автомобиль? Допустим у нас имеется изображение собаки. Для модели задача ставится так: Из картинки получить предсказание распределения вероятностей принадлежности к классам [0, 1, 0, 0]. Так как ничего идеального не бывает, то от обученной модели мы получим что-то похожее на [0.001, 0.99, 0.005, 0.004]. Заметьте сумма собирается в 1, как и положено свойству вероятности. А от необученной модели получим вообще всё что угодно. Например, [0.4, 0.15, 0.33, 0.12].

Давайте представим, что вот наша картинка собачки. Это конечно не так, но для модели разницы нет!

input_x = np.array([[ 0.99197708, -0.77980023, -0.8391331 , -0.41970686,  0.72636492],
       [ 0.85901409, -0.22374584, -1.95850625, -0.81685145,  0.96359871],
       [-0.42707937, -0.50053309,  0.34049477,  0.62106931, -0.76039365],
       [ 0.34206742,  2.15131285,  0.80851759,  0.28673013,  0.84706839],
       [-1.70231094,  0.36473216,  0.33631525, -0.92515589, -2.57602677]]) # абсолютно случайный массив

А на выходе хотим получить вероятности принадлежности к классам собака и кошка, то есть:

target_x = [1.0, 0.0]

Возьмем простую 2-х слойную нейронную сетью и попробуем обучить модель на данную простую задачу. Хочу заметить размерность input_x -> (5, 5). В контексте линейного слоя это означает 5 векторов размерности 5. Так как у нас одна картинка, то нам нужно перевести матрицу из (5,5) в (1, 25). В NumPy, метод reshape() используется для изменения формы массива без изменения его данных. Это позволяет преобразовывать одномерные массивы в многомерные и обратно, а также изменять размерности, если это возможно. Но в PyTorch для этого есть класс Flatten. Поэтому давайте также напишем свой.

class Flatten:
    def __init__(self): 
        pass
    def __call__(self, x):
        return x.reshape(1, -1)

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

class ReLU:
    def __init__(self): 
        pass
    def __call__(self, x):
        self.x = x
        return np.maximum(0, x) # поэлементно сравнивает значение с нулём

Что ж, пора уже собирать полноценную нейронку!

linear1 = Linear(input_channels=25, output_channels=10, bias=True)
linear2 = Linear(input_channels=10, output_channels=2, bias=True)
flatten = Flatten()
relu = ReLU()

x_1 = flatten(input_x)
x_2 = linear1(x_1)
x_3 = relu(x_2)
x_4 = linear2(x_3)
>>> array([[-0.11521861, -0.09036317]])

Ура, мы получили вывод нашей модели. Теперь будем с ним работать. Заметим что это не очень похоже на распределение вероятностей - значения отрицательные и не суммируются в единицу. Значит нам нужно придумать, как конвертировать эти значения в понятные нам вероятности. Давайте сперва избавимся от отрицательных значений прибавив какое-нибудь число, например, 0.2. Тогда: [-0.11521861 + 0.2, -0.09036317 + 0.2] = [0.08478139, 0.10963683].

Потом поделим на их сумму, чтобы нормировать на 1, [0.08478139, 0.10963683] / (0.08478139 + 0.10963683) = [0.43607738, 0.56392262]

Круто! Теперь что-то осмысленное. Но давайте рассмотрим еще одну технику. Избавиться от отрицательных чисел можно, например, возведением в степень.

2 ^ (0.5) = 1.41, 2 ^ (-0.5) = 0.70

И вообще a ^ b > 0 для любого b и любого положительного a. Тогда в нашем случае: [2 ^ (-0.11521861), 2 ^ (-0.09036317)] = [0.9232424 , 0.93928628]]. Снова нормируем и получаем [0.49569299, 0.50430701] Круто! На практике же используют вместо 2, число e. Причина станет понятна чуть позже.

Итак, это операция называется softmax

И для нее также есть свой класс Softmax в PyTorch. Поэтому давайте реализуем его!

class Softmax():
    def __init__(self): 
        pass
    def __call__(self, z):
        return np.exp(z) / np.sum(np.exp(z), axis=1).reshape(-1, 1) # reshape(-1, 1) для корректности

softmax = Softmax()
output = softmax(x_4)
>>> array([[0.49378646, 0.50621354]])

Отлично! Теперь у нас имеется 2 матрицы target_x и output. И как я объяснял выше, нам нужно что-то сделать с весами нашей модели, чтобы приблизить output к target_x . Для этого мы используем алгоритм градиентного спуска. Я подозреваю, что вы знакомы с ним поэтому сразу перейду к вычислению loss-функции! Для задачи классификации мы используем loss-функцию - кросс-энтропия.

В PyTorch для нее также есть свой класс, а значит и мы попробуем его создать!

class CrossEntropyLoss:
    def __init__(self):
        self.predicted = None
        self.true = None

    def __call__(self, predicted, true):
        self.predicted = np.array(predicted, copy=True) 
        self.true = np.array(true, copy=True) # сделаем копию входных матриц для дальнейших вычислений
        # вычисляем значение лосс-функции прямо по формуле
        self.loss = -1 * np.sum(true * np.log(predicted + 1e-5), axis=1) # добавим 1e-5, чтобы случайно не взять log(0)
        return self
        
loss_fn = CrossEntropyLoss()
loss = loss_fn(output, target_x)
loss.loss
>>> array([0.68077693])

Круто! С каждым разом мы всё ближе и ближе к настоящим большим нейронкам!

Обучение

Итак, мы получили значение loss-функции - некоторое число. Дальше, мы знаем что нужно прибегнуть к градиентному спуску. То есть много раз дифференцировать наше значение, чтобы получить необходимые изменения весов. Или грубо говоря, нам нужно реализовать метод .backward() у класса CrossEntropyLoss. Давайте реализуем самый самый простой алгоритм градиентного спуска. Для этого нам понадобится реализовать .backward() в каждом классе. Приступим!

Линейный слой

Якобиан произведения двух матриц можно выразить следующим образом:
Если C = AB , то можно записать производную функции потерь по матрице следующим образом:

$\frac{dL}{dA} = \frac{dL}{dC} \cdot \frac{dC}{dA}$

Для вычисления производной по матрице (где C = AB ):

$\frac{dC}{dA} = B^T$

Таким образом:

$\frac{dL}{dA} = \frac{dL}{dC} \cdot B^T$

Аналогично для производной по матрице :

$\frac{dL}{dB} = A^T \cdot \frac{dL}{dC}$

class Linear:
    def backward(self, input_matrix):
        x_gradient = input_matrix @ self.weight.T
        self.weight_gradient = self.x.T @ input_matrix
        self.bias_gradient = input_matrix.mean(axis=0)
        return x_gradient # возвращаем чтобы посчитать градиенты через следующие слои

Функция активации

Для ReLu несложно, там где положительные значения возвращаем само значение, а там где были отрицательные возвращаем 0!

class ReLU:
    def backward(self, input_matrix):
        return (self.x > 0) * input_matrix

Кросс-Энтропия

Используя формулу для производной кросс-энтропии (доказательство), в итоге получаем

back1 = loss.predicted - loss.true
back2 = linear2.backward(back1)
back3 = relu.backward(back2)
back4 = linear1.backward(back3)
back4.shape
>>> (1, 25)

Обновление весов

Как же нам теперь обновить веса? Очень просто! Сделаем вручную градиентный спуск для каждого слоя.

linear1.weight = linear1.weight - 0.001*linear1.weight_gradient
linear1.bias = linear1.bias - 0.001*linear1.bias_gradient

linear2.weight = linear2.weight - 0.001*linear2.weight_gradient
linear2.bias = linear2.bias - 0.001*linear2.bias_gradient

Попробуем уже обучить нашу нейронку!

for i in range(100):
    x_1 = flatten(input_x)
    x_2 = linear1(x_1)
    x_3 = relu(x_2)
    x_4 = linear2(x_3)
    x_5 = softmax(x_4)
    loss = loss_fn(x_5, target_x)
    back1 = loss.predicted - loss.true
    back2 = linear2.backward(back1)
    back3 = relu.backward(back2)
    back4 = linear1.backward(back3)
    
    lr = 0.01
    linear1.weight = linear1.weight - lr*linear1.weight_gradient
    linear1.bias = linear1.bias - lr*linear1.bias_gradient
    
    linear2.weight = linear2.weight - lr*linear2.weight_gradient
    linear2.bias = linear2.bias - lr*linear2.bias_gradient

    if (i % 20) == 0:
        print(loss.loss, i)
        
>>>[0.60362514] 0
[0.11279846] 20
[0.05461399] 40
[0.03388313] 60
[0.02385492] 80

Ура, у нас всё получилось!!!

Это очень круто, потому что мы убедились что у нас все работает, значение loss-фукнции уменьшается практически до 0. Но хочется пойти дальше! Что мы можем улучшить, дополнить, изменить? На самом деле очень много. Давайте, посмотрим на код, очевидно, что нам не хочется каждый раз руками контролировать переток градиентов, а хочется, чтобы за нас сделала это программа! Может попробуем засунуть все вычисления в метод .backward() класса CrossEntropyLoss? Но тут возникает небольшая проблема. Вы знаете, что в градиентном спуске нам нужно дифференцировать полученное значение loss-функции по разным входам, выходам и параметрам сети. Как вы могли заметить мы получили просто число, а через какие слои и каким образом проходило это значение, непосредственно из этого значения мы получить не можем! Значит нам каким-то образом нужно передать информацию о всех слоях нашей нейросети и о том как изменяются данные в этих слоях! (В дальнейшем эта проблема будет решаться с помощью вычислительного графа, но пока обойдемся без него!)

Итак, предлагаю несколько нововведений! Давайте теперь определять нашу модель внутри некоторого класса, который будет наследоваться от созданного нами ранее родительского класса Module. Например

class SimpleNet(Module):
    def __init__(self):
        super().__init__()
        self.linear1 = Linear(input_channels=25, output_channels=10, bias=True)
        self.linear2 = Linear(input_channels=10, output_channels=2, bias=True)
        self.flatten = Flatten()
        self.relu = ReLU()
        self.softmax = Softmax()

    def forward(self, x):
        x_1 = self.flatten(x)
        x_2 = self.linear1(x_1)
        x_3 = self.relu(x_2)
        x_4 = self.linear2(x_3)
        x_5 = self.softmax(x_4)
        return x_5
model = SimpleNet()
# теперь весь вызов модели будет в одну строчку
output = model(input_x)

Наша задача с помощью loss.backward() посчитать все-все градиенты, при этом передавать какую-то информацию о модели в этот метод мы не можем, потому что это будет нечестно так как в PyTorch мы ничего не передаем. Нам каким-то образом нужно сделать метку всем слоям одной модели. Например, если линейный слой был создан внутри первой модели то у него будет первый флажок, а если он был создан внутри второй модели то второй флажок, отличный от первого. Так по флажкам мы сможем определить какие слои какой модели принадлежат и в каком порядке они используются, а значит корректно посчитать градиенты.

Введем глобальную переменную Parameter = None - именно она будет нашим флажком! Воспользуемся замыканием

def ParameterObj():
    class Parameter:
        layers = []
        calling = dict()
        def __init__(self, info): # info(ссылка на объект, *веса слоя)
            Parameter.layers.append(info[0])
            Parameter.calling[info[0]] = info[1:]
    return Parameter
    

class Module:
    def __init__(self):
        self._constructor_Parameter = ParameterObj()
        global Parameter
        Parameter = self._constructor_Parameter

Когда мы будем создавать новую модель - будет создаваться внутренний класс. Но этот класс будет разным для разных моделей (это называется замыкание), поэтому мы сможем отличать одну модель от другой. Во время создания модели мы ссылаем глобальную переменную Parameter на новый класс, и слои которые будут создаваться в этой модели, при обращении к глобальной переменной будут получать ссылки именно на класс принадлежащий этой модели!

Теперь с учетом этого факта нам нужно немного поменять наши созданные классы

class Linear:
    def __init__(self, input_channels: int, output_channels: int, bias = True):
        self.input_channels = input_channels
        self.output_channels = output_channels
        self.bias = bias
        self.backward_list = [] # храним градиенты для параметров слоя
        global Parameter
        if bias:
            Parameter([self, np.random.uniform(- 0.5, 0.5, size=(self.input_channels, self.output_channels)), np.random.uniform(- 0.5, 0.5, size=self.output_channels)])
        else:
            Parameter([self, np.random.uniform(- 0.5, 0.5, size=(self.input_channels, self.output_channels)),np.zeros(self.output_channels)])

    def __call__(self, x):
        result = x @ Parameter.calling[self][0] + Parameter.calling[self][1]
        return result

Отлично, теперь у нас все слои модели связаны при помощи глобальной переменной и её атрибутов layers и calling. Первое хранит ссылки на объекты слоёв. Второе хранит сами параметры этих слоёв!

Начнем писать метод .backward() для класса CrossEntropyLoss

def backward(self):
        # производная кросс-энтропии
        loss = self.predicted - self.true
        # Итерируем по каждому слою в обратном порядке, благодаря тому, что мы всё сохранили в Parameter.layers
        for index, layer in enumerate(Parameter.layers[::-1]):
            if type(layer).__name__ == 'Linear':
                changes_w = (layer.x.T @ loss) / loss.shape[0] # нормировка на loss.shape[0] нужна, так как величина изменений зависит от размера батча
                if layer.bias:
                    changes_b = (np.sum(loss) / loss.shape[0])
                else:
                    changes_b = 0 # Не меняем ничего, если изначально не использовали bias
                layer.backward_list = [changes_w, changes_b]
                # Считаем градиент для протекания к следующим слоям
                loss = loss @ Parameter.calling[layer][0].T
                
            elif type(layer).__name__ == 'ReLU':
                loss = layer.backward(loss)

Ничего сложного, давайте проверим работу!

loss_fn = CrossEntropyLoss()
model = SimpleNet()

for i in range(100):
    output = model(input_x)
    loss = loss_fn(output, target_x)
    loss.backward()
    lr = 0.01

    for layer in (model.linear1, model.linear2):
        weight, bias = model._constructor_Parameter.calling[layer]
        weight_gradient, bias_gradient = layer.backward_list[0], layer.backward_list[1]
        new_weight = weight - lr * weight_gradient
        new_bias = bias - lr * bias_gradient
        model._constructor_Parameter.calling[layer] = [new_weight, new_bias] # теперь при вызове будем получать обновленные параметры слоя
        
    if (i % 20) == 0:
        print(loss.loss, i)

>>> [0.68439962] 0
[0.05159733] 20
[0.02452264] 40
[0.01580639] 60
[0.01160133] 80

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

Для этого создадим новый класс который будет ответственным за изменение параметров модели.

Stochastic gradient descent

class SGD:
    def __init__(self, model, learning_rate):
        self.model = model
        self.lr = learning_rate

    def step(self):
        # с помощью self.model._constructor_Parameter получаем ссылку на нужный класс
        for index, layer in enumerate(self.model._constructor_Parameter.layers[::-1]):
            if type(layer).__name__  == 'Linear':
                weight, bias = self.model._constructor_Parameter.calling[layer]
                weight_gradient, bias_gradient = layer.backward_list
                new_weight = weight - self.lr * weight_gradient
                new_bias = bias - self.lr * bias_gradient
                self.model._constructor_Parameter.calling[layer] = [new_weight, new_bias]

Напомню, как мы инициализируем оптимизатор.

optimizer = optim.SGD(model.parameters(), lr=0.01)

Мы передаем параметры модели, а у нас мы передаем саму модель. Предлагаю костыль, который мы потом как-нибудь исправим.

class Module:
    def parameters(self):
        return self

Теперь собираем!

loss_fn = CrossEntropyLoss()
model = SimpleNet()
optim = SGD(model.parameters(), learning_rate = 0.01)

for i in range(100):
    output = model(input_x)
    loss = loss_fn(output, target_x)
    loss.backward()
    optim.step()
    
    if (i % 20) == 0:
        print(loss.loss, i)
        
>>> [0.18273859] 0
[0.04809317] 20
[0.02612311] 40
[0.01764624] 60
[0.01320964] 80

Заключительная косметическая правка. В PyTorch в loss-функцию мы передаем не вероятности, а логиты.

То есть нам не нужно применять softmax к выходам модели самостоятельно, за нас должна сделать это функция, поэтому

class SimpleNet(Module):

    def forward(self, x):
        x_1 = self.flatten(x)
        x_2 = self.linear1(x_1)
        x_3 = self.relu(x_2)
        x_4 = self.linear2(x_3)
        return x_4

class CrossEntropyLoss:

    def __call__(self, logits, true):
        predicted = np.exp(logits) / np.sum(np.exp(logits), axis=1).reshape(-1, 1) # softmax
        self.predicted = np.array(predicted, copy=True) # сделаем копию входных матрицы для дальнейших вычислений
        self.true = np.array(true, copy=True) # сделаем копию входных матрицы для дальнейших вычислений
        number_of_classes = predicted.shape[1] # получим количество классов, в нашем случае 2
        self.true = np.array(true, copy=True)
        # вычисляем значение лосс-функции прямо по формуле
        self.loss = -1 * np.sum(true * np.log(predicted + 1e-5), axis=1)
        return self

Вот и на этом первая часть подошла к концу. Полный код можно найти тут

Мы построили минимальную обёртку для NumPy, которая уже неплохо имитирует работу PyTorch. Но работы ещё много.

Во второй части я планирую:

- добавить CNN, RNN, BatchNorm, Dropout, MaxPool, MinPool
- реализовать RMSProp, NaG, Adam
- добавить регуляризацию в loss-функцию
- добавить новые функции активации
- написать DataLoader

В третьей части я планирую:

- представить аналог pytorch.tensor()
- перевести все вычисления на динамический вычислительный граф
- добавить Embedding, LayerNorm, ModuleList
- провести рефакторинг библиотеки
- добавить перенос вычислений на gpu
- написать на библиотеке gpt2-1.5В и запустить его

Первая версия библиотеки

Вторая версия библиотеки

GPT-2 на этой библиотеке

Пишем свой PyTorch на NumPy. Часть 1