PyTorch: Learn the Basics, часть 3 — модель и autograd

Третья часть перевода курса Learn the Basics (исходники — pytorch/tutorials, BSD-3-Clause): разделы «Build the Neural Network» и «Automatic Differentiation with torch.autograd». Предыдущая часть — Dataset, DataLoader и преобразования, следующая — цикл обучения и сохранение.

Как читать вместе с курсом ChinaAI

Если вы прошли модуль micrograd, раздел об autograd будет узнаваемым: grad_fn — это то же, что _backward у класса Value, а .grad накапливается точно так же. Устройство слоя nn.Linear и функции ReLU объясняет урок «Нейрон, слой, MLP», цепное правило — урок «Цепное правило и вычислительный граф».

4. Сборка нейросети

Нейросети состоят из слоёв (модулей), которые выполняют операции над данными. Пространство имён torch.nn даёт все строительные блоки, нужные, чтобы собрать свою нейросеть. Каждый модуль в PyTorch наследует nn.Module. Нейросеть сама по себе — модуль, который состоит из других модулей (слоёв). Такая вложенная структура позволяет легко строить сложные архитектуры и управлять ими.

В следующих подразделах мы соберём нейросеть, которая классифицирует изображения из набора FashionMNIST.

import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

Выбираем устройство для обучения

Мы хотим обучать модель на ускорителе — CUDA, MPS, MTIA или XPU. Если ускоритель доступен, используем его, иначе — CPU.

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")

Вывод:

Using cuda device

Определяем класс

Мы задаём нейросеть как подкласс nn.Module и создаём её слои в __init__. Каждый подкласс nn.Module описывает операции над входными данными в методе forward.

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

Создаём экземпляр NeuralNetwork, переносим его на device и печатаем его структуру.

model = NeuralNetwork().to(device)
print(model)

Вывод:

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

Чтобы воспользоваться моделью, мы передаём ей входные данные. При этом выполняется метод forward модели вместе с некоторыми фоновыми операциями. Не вызывайте model.forward() напрямую!

Вызов модели на входе возвращает двумерный тензор: измерение dim=0 соответствует отдельным выходам (по одному на каждый пример), а dim=1 — десяти «сырым» предсказанным значениям для каждого класса. Вероятности предсказаний мы получаем, пропуская результат через экземпляр модуля nn.Softmax.

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

Вывод:

Predicted class: tensor([0], device='cuda:0')

(Примечание переводчика: у тензора logits здесь форма [1, 10] — один пример, десять классов. Сеть ещё не обучена, поэтому предсказанный класс случаен.)

---

Слои модели

Разберём по слоям модель для FashionMNIST. Для наглядности возьмём мини-батч из трёх изображений размером 28×28 и посмотрим, что с ним происходит при проходе через сеть.

input_image = torch.rand(3,28,28)
print(input_image.size())

Вывод:

torch.Size([3, 28, 28])

nn.Flatten

Мы создаём слой nn.Flatten, чтобы превратить каждое двумерное изображение 28×28 в непрерывный массив из 784 значений пикселей (измерение мини-батча, dim=0, сохраняется).

flatten = nn.Flatten()
flat_image = flatten(input_image)
print(flat_image.size())

Вывод:

torch.Size([3, 784])

nn.Linear

Линейный слой — это модуль, который применяет к входу линейное преобразование с помощью хранящихся в нём весов и смещений.

layer1 = nn.Linear(in_features=28*28, out_features=20)
hidden1 = layer1(flat_image)
print(hidden1.size())

Вывод:

torch.Size([3, 20])

(Примечание переводчика: внутри слой считает $y = xW^\top + b$, где у матрицы $W$ форма $20 \times 784$, а у вектора смещений $b$ — 20 чисел. Всего у такого слоя $20 \cdot 784 + 20 = 15\,700$ обучаемых параметров.)

nn.ReLU

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

В этой модели мы ставим nn.ReLU между линейными слоями, но есть и другие функции активации, которые вносят нелинейность в модель.

print(f"Before ReLU: {hidden1}\n\n")
hidden1 = nn.ReLU()(hidden1)
print(f"After ReLU: {hidden1}")

Вывод (первая строка каждого тензора):

Before ReLU: tensor([[-0.3729,  0.1928, -0.0521,  0.0542,  0.2942, -0.1856,  0.6333, -0.0815,
         -0.0360,  0.5083, -0.0697,  0.1799,  0.0639,  0.0892, -0.0324,  0.0707,
          0.6500,  0.6037,  0.5211, -0.1809],
        ...], grad_fn=<AddmmBackward0>)


After ReLU: tensor([[0.0000, 0.1928, 0.0000, 0.0542, 0.2942, 0.0000, 0.6333, 0.0000, 0.0000,
         0.5083, 0.0000, 0.1799, 0.0639, 0.0892, 0.0000, 0.0707, 0.6500, 0.6037,
         0.5211, 0.0000],
        ...], grad_fn=<ReluBackward0>)

(Примечание переводчика: ReLU заменяет отрицательные числа нулями, а положительные оставляет как есть: −0.3729 стало 0.0000, а 0.1928 не изменилось.)

nn.Sequential

nn.Sequential — упорядоченный контейнер модулей. Данные проходят через все модули в том порядке, в котором они заданы. С помощью таких последовательных контейнеров можно быстро собрать сеть, например seq_modules.

seq_modules = nn.Sequential(
    flatten,
    layer1,
    nn.ReLU(),
    nn.Linear(20, 10)
)
input_image = torch.rand(3,28,28)
logits = seq_modules(input_image)

nn.Softmax

Последний линейный слой нейросети возвращает логиты (logits) — «сырые» значения в диапазоне $[-\infty, \infty]$, — которые передаются в модуль nn.Softmax. Логиты масштабируются в значения из $[0, 1]$ — вероятности каждого класса по мнению модели. Параметр dim указывает измерение, вдоль которого значения должны в сумме давать 1.

softmax = nn.Softmax(dim=1)
pred_probab = softmax(logits)

(Примечание переводчика: формула softmax и её связь с функцией потерь разобраны в уроке «Правдоподобие, softmax и cross-entropy».)

Параметры модели

Многие слои внутри нейросети параметризованы: у них есть веса и смещения, которые оптимизируются при обучении. Наследование от nn.Module автоматически отслеживает все поля, заданные внутри объекта модели, и делает все параметры доступными через методы модели parameters() и named_parameters().

В этом примере мы перебираем все параметры и печатаем размер каждого и первые значения.

print(f"Model structure: {model}\n\n")

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \n")

Вывод (после структуры модели, которую мы уже видели выше):

Layer: linear_relu_stack.0.weight | Size: torch.Size([512, 784]) | Values : tensor([[-0.0142,  0.0019,  0.0209,  ...,  0.0028, -0.0299,  0.0136],
        [-0.0104, -0.0318, -0.0339,  ...,  0.0230, -0.0129,  0.0068]],
       device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.0.bias | Size: torch.Size([512]) | Values : tensor([-0.0064,  0.0003], device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.2.weight | Size: torch.Size([512, 512]) | Values : tensor([[ 0.0339, -0.0120, -0.0293,  ...,  0.0294,  0.0118, -0.0192],
        [-0.0193,  0.0286, -0.0279,  ...,  0.0342,  0.0013, -0.0320]],
       device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.2.bias | Size: torch.Size([512]) | Values : tensor([-0.0030,  0.0291], device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.4.weight | Size: torch.Size([10, 512]) | Values : tensor([[-0.0156,  0.0416,  0.0267,  ...,  0.0242, -0.0294,  0.0353],
        [ 0.0196, -0.0432,  0.0250,  ...,  0.0259,  0.0265,  0.0191]],
       device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.4.bias | Size: torch.Size([10]) | Values : tensor([ 0.0317, -0.0178], device='cuda:0', grad_fn=<SliceBackward0>)

(Примечание переводчика: сложите размеры — $784 \cdot 512 + 512 + 512 \cdot 512 + 512 + 512 \cdot 10 + 10 = 669\,706$ параметров. Проверить можно одной строкой: sum(p.numel() for p in model.parameters()).)

Что почитать дальше

• API torch.nn

---

5. Автоматическое дифференцирование с torch.autograd

При обучении нейросетей чаще всего используют алгоритм обратного распространения ошибки (back propagation). В нём параметры (веса модели) подстраиваются в соответствии с градиентом функции потерь по этому параметру.

Чтобы вычислять эти градиенты, в PyTorch есть встроенный механизм дифференцирования — torch.autograd (автоматическое дифференцированиеавтоматическое дифференцирование (автоград)EN automatic differentiation (autograd)Программа записывает операции в граф и сама считает производные по цепному правилу. Для loss → параметры используют обратный режим (reverse mode) — это и есть backprop.). Он умеет автоматически вычислять градиент для любого вычислительного графа.

Возьмём простейшую однослойную нейросеть со входом x, параметрами w и b и некоторой функцией потерь. В PyTorch её можно задать так:

import torch

x = torch.ones(5)  # входной тензор
y = torch.zeros(3)  # ожидаемый выход
w = torch.randn(5, 3, requires_grad=True)
b = torch.randn(3, requires_grad=True)
z = torch.matmul(x, w)+b
loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)

Тензоры, функции и вычислительный граф

Этот код задаёт такой вычислительный граф (computational graph):

flowchart LR
x["x"] --> mul(("*"))
w["w (параметр)"] --> mul
mul --> add(("+"))
b["b (параметр)"] --> add
add --> z["z"]
z --> ce["CE"]
y["y"] --> ce
ce --> loss["loss"]

Вычислительный граф из примера: z = x·w + b, затем бинарная перекрёстная энтропия с y (перерисовано переводчиком по рисунку comp-graph.png из оригинала)

В этой сети w и b — параметры, которые нужно оптимизировать. Поэтому нам нужно уметь вычислять градиенты функции потерь по этим переменным. Для этого мы устанавливаем у этих тензоров свойство requires_grad.

Заметка

Значение requires_grad можно задать при создании тензора, а можно позже — методом x.requires_grad_(True).

Функция, которую мы применяем к тензорам, чтобы построить вычислительный граф, на самом деле — объект класса Function. Этот объект знает, как вычислить функцию в прямом направлении и как вычислить её производную на шаге обратного распространения. Ссылка на функцию обратного распространения хранится в свойстве тензора grad_fn. Подробнее о Function — в документации.

print(f"Gradient function for z = {z.grad_fn}")
print(f"Gradient function for loss = {loss.grad_fn}")

Вывод:

Gradient function for z = <AddBackward0 object at 0x7f859f581180>
Gradient function for loss = <BinaryCrossEntropyWithLogitsBackward0 object at 0x7f859f582920>

Вычисление градиентов

Чтобы оптимизировать веса параметров нейросети, нужно вычислить производные функции потерь по параметрам, то есть нам нужны $\frac{\partial\, loss}{\partial w}$ и $\frac{\partial\, loss}{\partial b}$ при некоторых фиксированных значениях x и y. Чтобы вычислить эти производные, мы вызываем loss.backward(), а затем берём значения из w.grad и b.grad:

loss.backward()
print(w.grad)
print(b.grad)

Вывод:

tensor([[0.2752, 0.0590, 0.1873],
        [0.2752, 0.0590, 0.1873],
        [0.2752, 0.0590, 0.1873],
        [0.2752, 0.0590, 0.1873],
        [0.2752, 0.0590, 0.1873]])
tensor([0.2752, 0.0590, 0.1873])

(Примечание переводчика: почему все строки w.grad одинаковы и совпадают с b.grad? Вход x — пять единиц, поэтому $\partial z_j / \partial w_{ij} = x_i = 1$, и градиент по каждому весу столбца $j$ равен градиенту по $z_j$ — то есть по $b_j$. Числа у вас будут другими: w и b случайные.)

Заметка

• Свойство grad можно получить только у листовых вершин вычислительного графа, у которых requires_grad равно True. Для всех остальных вершин графа градиенты недоступны.
• Ради производительности вычислять градиенты с помощью backward на одном и том же графе можно только один раз. Если нужно несколько вызовов backward на одном графе, передайте в вызов backward аргумент retain_graph=True.

Отключение отслеживания градиентов

По умолчанию все тензоры с requires_grad=True отслеживают историю своих вычислений и поддерживают вычисление градиента. Но бывают случаи, когда это не нужно: например, когда модель уже обучена и мы просто хотим применить её к входным данным, то есть выполнить только прямой проход по сети. Остановить отслеживание вычислений можно, обернув код в блок torch.no_grad():

z = torch.matmul(x, w)+b
print(z.requires_grad)

with torch.no_grad():
    z = torch.matmul(x, w)+b
print(z.requires_grad)

Вывод:

True
False

Другой способ добиться того же — метод detach() у тензора:

z = torch.matmul(x, w)+b
z_det = z.detach()
print(z_det.requires_grad)

Вывод:

False

Причины отключить отслеживание градиентов:

• чтобы пометить некоторые параметры нейросети как замороженные (frozen parameters);
• чтобы ускорить вычисления, когда нужен только прямой проход: вычисления над тензорами, которые не отслеживают градиенты, эффективнее.

(Примечание переводчика: заморозка параметров — основа дообучения больших моделей. В LoRA веса исходной модели заморожены, а обучаются только маленькие добавочные матрицы.)

---

Подробнее о вычислительных графах

Концептуально autograd хранит запись данных (тензоров) и всех выполненных операций (вместе с получившимися новыми тензорами) в направленном ациклическом графе (DAG), который состоит из объектов Function. В этом графе листья — входные тензоры, а корни — выходные. Проходя граф от корней к листьям, можно автоматически вычислить градиенты по цепному правилу.

При прямом проходе autograd делает одновременно две вещи:

• выполняет запрошенную операцию и вычисляет тензор-результат;
• сохраняет в графе функцию градиента этой операции.

Обратный проход запускается, когда у корня графа вызывают .backward(). Тогда autograd:

• вычисляет градиенты с помощью каждой .grad_fn;
• накапливает их в атрибуте .grad соответствующего тензора;
• по цепному правилу распространяет градиенты до самых листовых тензоров.

В PyTorch графы динамические

Важно: граф каждый раз строится заново; после каждого вызова .backward() autograd начинает заполнять новый граф. Именно это позволяет использовать в модели операторы управления потоком (условия, циклы): при необходимости на каждой итерации можно менять форму, размер и сами операции.

Дополнительно: градиенты тензоров и произведения с якобианом

Во многих случаях функция потерь скалярная, и нужно вычислить градиент по некоторым параметрам. Но бывает, что выход функции — произвольный тензор. Тогда PyTorch позволяет вычислить так называемое произведение с якобианом (Jacobian product), а не сам градиент.

Для векторной функции $\vec{y}=f(\vec{x})$, где $\vec{x}=\langle x_1,\dots,x_n\rangle$ и $\vec{y}=\langle y_1,\dots,y_m\rangle$, градиент $\vec{y}$ по $\vec{x}$ задаётся матрицей Якоби:

$$J=\left(\begin{array}{ccc} \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{1}}{\partial x_{n}}\\ \vdots & \ddots & \vdots\\ \frac{\partial y_{m}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}} \end{array}\right)$$

Вместо самой матрицы Якоби PyTorch позволяет вычислить произведение с якобианом $v^T\cdot J$ для заданного вектора $v=(v_1 \dots v_m)$. Для этого нужно вызвать backward с $v$ в качестве аргумента. Размер $v$ должен совпадать с размером исходного тензора, по которому мы хотим вычислить произведение:

inp = torch.eye(4, 5, requires_grad=True)
out = (inp+1).pow(2).t()
out.backward(torch.ones_like(out), retain_graph=True)
print(f"First call\n{inp.grad}")
out.backward(torch.ones_like(out), retain_graph=True)
print(f"\nSecond call\n{inp.grad}")
inp.grad.zero_()
out.backward(torch.ones_like(out), retain_graph=True)
print(f"\nCall after zeroing gradients\n{inp.grad}")

Вывод:

First call
tensor([[4., 2., 2., 2., 2.],
        [2., 4., 2., 2., 2.],
        [2., 2., 4., 2., 2.],
        [2., 2., 2., 4., 2.]])

Second call
tensor([[8., 4., 4., 4., 4.],
        [4., 8., 4., 4., 4.],
        [4., 4., 8., 4., 4.],
        [4., 4., 4., 8., 4.]])

Call after zeroing gradients
tensor([[4., 2., 2., 2., 2.],
        [2., 4., 2., 2., 2.],
        [2., 2., 4., 2., 2.],
        [2., 2., 2., 4., 2.]])

Заметьте: когда мы второй раз вызываем backward с тем же аргументом, значение градиента получается другим. Так происходит потому, что при обратном распространении PyTorch накапливает градиенты: вычисленные значения прибавляются к свойству grad всех листовых вершин графа. Чтобы получить правильные градиенты, свойство grad нужно предварительно обнулить. При реальном обучении это делает за нас оптимизатор.

(Примечание переводчика: проверим числа. Производная $(x+1)^2$ равна $2(x+1)$. На диагонали torch.eye стоят единицы: $2 \cdot (1+1) = 4$; вне диагонали нули: $2 \cdot (0+1) = 2$. Второй вызов прибавил то же самое — получилось 8 и 4.)

Заметка

Раньше мы вызывали backward() без параметров. По сути это то же, что backward(torch.tensor(1.0)) — удобный способ вычислить градиенты для скалярной функции, например функции потерь при обучении нейросети.

Что почитать дальше

• Autograd Mechanics

---

Словарик

EN	RU
layer, module	слой, модуль
linear layer	линейный (полносвязный) слой
activation function	функция активации
logits	логиты
weights, biases	веса, смещения
automatic differentiation	автоматическое дифференцирование
computational graph	вычислительный граф
leaf node	листовая вершина
gradient accumulation	накопление градиентов
frozen parameters	замороженные параметры
Jacobian matrix	матрица Якоби (якобиан)

Дальше: часть 4 — цикл обучения, сохранение и загрузка модели.