PyTorch: autograd и автоматическое дифференцирование

В первом уроке мы вывели градиенты линейной регрессии вручную — всего для двух параметров. В современных нейросетях миллионы параметров, в больших языковых моделях — миллиарды; выводить каждую производную на бумаге невозможно. Механизм autograd в PyTorch делает это автоматически: он запоминает каждую операцию с тензорами, а затем одной командой вычисляет точные производные по всем параметрам. Это не приближённые численные разности и не символьная алгебра — это автоматическое дифференцирование.

Вычислительный граф

Каждая операция с тензором, у которого requires_grad=True, добавляет узел в ориентированный ациклический граф (DAG). Листья графа — входы и параметры, корень — функция потерь. Во время прямого прохода (forward pass) граф строится и вычисляются значения; во время обратного прохода (backward pass) autograd идёт от корня к листьям и в каждом узле применяет цепное правило. В нейросетях этот процесс называется обратным распространением ошибки (backpropagation).

В PyTorch граф динамический (define-by-run): он строится заново при каждом прямом проходе. Поэтому внутри модели можно свободно использовать обычные if и for из Python — граф будет соответствовать тому пути, который реально выполнился.

requires_grad, backward() и .grad

import torch

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2 + 2 * x + 1
print(y)
y.backward()
print(x.grad)

tensor(16., grad_fn=<AddBackward0>)
tensor(8.)

• requires_grad=True говорит PyTorch: «по этому тензору понадобится производная, запоминай операции». Параметры (веса nn.Linear и т. д.) делают это автоматически.
• grad_fn показывает, какая операция создала тензор; это его узел в графе.
• y.backward() вычисляет dy/dx и записывает результат в .grad листа: для y = x² + 2x + 1 при x = 3 dy/dx = 2x + 2 = 8.

Если параметр — вектор, то .grad — вектор той же формы: по частной производной на каждый элемент. Например, для f = ∑ vᵢ² имеем ∂f/∂vᵢ = 2vᵢ:

import torch

v = torch.tensor([1.0, -2.0, 3.0], requires_grad=True)
f = (v ** 2).sum()
f.backward()
print(v.grad)

tensor([ 2., -4.,  6.])

Флаг requires_grad можно менять и потом: p.requires_grad_(False) «замораживает» параметр — его градиент не вычисляется, и при обучении он не меняется. Этим пользуются при переносе обучения (transfer learning): большинство слоёв предобученной сети замораживают, а под новую задачу дообучают только последний слой. Это мы увидим в уроке о свёрточных сетях.

Цепное правило

import torch

w = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
x, y = torch.tensor(3.0), torch.tensor(10.0)
y_hat = w * x + b
loss = (y_hat - y) ** 2
print(type(loss.grad_fn).__name__, type(y_hat.grad_fn).__name__)
loss.backward()
print(loss.item(), w.grad.item(), b.grad.item())

PowBackward0 AddBackward0
9.0 -18.0 -6.0

import torch

w = torch.tensor(0.0, requires_grad=True)
x = torch.tensor(2.0)
y = torch.sigmoid(w * x)
y.backward()
print(y.item(), w.grad.item())

0.5 0.5

Сила обратного прохода — в эффективности: один вызов backward() вычисляет градиент потерь по всем параметрам, а стоит это примерно столько же, сколько прямой проход (с небольшим постоянным множителем). Если бы мы считали каждый параметр отдельно, для сети с миллионом параметров понадобился бы миллион прямых проходов.

Градиенты накапливаются: обнуление и torch.no_grad()

.backward() не заменяет старый градиент, а прибавляет к нему. Ниже производная y = x³ при x = 2 равна 3x² = 12, но после трёх вызовов .grad равен 36:

import torch

x = torch.tensor(2.0, requires_grad=True)
for i in range(3):
    y = x ** 3
    y.backward()
    print(x.grad)
x.grad.zero_()
print(x.grad)
with torch.no_grad():
    z = x * 2
print(z.requires_grad, x.detach().requires_grad)

tensor(12.)
tensor(24.)
tensor(36.)
tensor(0.)
False False

• x.grad.zero_() обнуляет градиент на месте (подчёркивание _ в конце означает операцию на месте). В следующих уроках это будет делать optimizer.zero_grad().
• В блоке with torch.no_grad(): граф не строится: его используют при обновлении параметров и при оценке модели (инференсе), это экономит память и время.
• .detach() возвращает тензор с теми же данными, но отсоединённый от графа — например, чтобы вывести значение или преобразовать его в NumPy.

Градиентный спуск с autograd

Теперь соберём всё вместе и подгоним линейную модель к 20 зашумлённым точкам прямой y = 2x + 1. На этот раз мы не пишем формулы производных — autograd вычисляет их на каждом шаге. Пять шагов этого цикла — основа всего глубокого обучения: прямой проход → потери → backward → обновление внутри no_grad → обнуление градиентов.

import torch

torch.manual_seed(42)
X = torch.linspace(0, 1, 20)
Y = 2 * X + 1 + 0.1 * torch.randn(20)
w = torch.zeros(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.5
for step in range(201):
    loss = ((w * X + b - Y) ** 2).mean()
    loss.backward()
    with torch.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad
    w.grad.zero_()
    b.grad.zero_()
    if step % 50 == 0:
        print(f'step {step:3d}  loss {loss.item():.4f}')
print(f'w = {w.item():.3f}, b = {b.item():.3f}')

step   0  loss 4.4302
step  50  loss 0.0085
step 100  loss 0.0084
step 150  loss 0.0084
step 200  loss 0.0084
w = 1.907, b = 1.068

Хорошо запомни этот цикл: в следующих уроках мы его сократим, но суть не изменится. Метод step() оптимизатора torch.optim.SGD выполняет для каждого параметра именно p -= lr * p.grad внутри no_grad, а zero_grad() очищает градиенты. Слои вроде nn.Linear сами создают w и b и сами включают requires_grad. То есть «волшебный» высокоуровневый код — лишь удобная обёртка над этими пятью строками.