PyTorch: цикл обучения модели

Теперь у нас есть все детали: тензоры, autograd и nn.Module. Соберём их в стандартный рецепт, который специалисты пишут каждый день: данные → DataLoader → модель → потери → оптимизатор → цикл обучения → оценка → сохранение. Любопытно, что этот каркас почти одинаков и для крошечного классификатора, и для языковой модели с миллиардами параметров — меняется лишь размер данных и модели.

Dataset и DataLoader

Dataset — объект, который ведёт себя как список примеров: __len__ возвращает число примеров, а __getitem__(i) — i-ю пару (x, y). Если данные уже в тензорах, хватит готового TensorDataset; для чтения из файлов (изображений, текстов) пишут собственный класс. DataLoader собирает из Dataset пакеты, при необходимости перемешивает примеры каждую эпоху (shuffle=True) и умеет загружать данные параллельно через num_workers.

import torch
from torch.utils.data import Dataset, DataLoader

class SquaresDataset(Dataset):
    def __init__(self, n):
        self.x = torch.arange(n, dtype=torch.float32)

    def __len__(self):
        return len(self.x)

    def __getitem__(self, i):
        return self.x[i], self.x[i] ** 2

ds = SquaresDataset(10)
print(len(ds), ds[3])
loader = DataLoader(ds, batch_size=4)
for xb, yb in loader:
    print(xb.tolist(), yb.tolist())

10 (tensor(3.), tensor(9.))
[0.0, 1.0, 2.0, 3.0] [0.0, 1.0, 4.0, 9.0]
[4.0, 5.0, 6.0, 7.0] [16.0, 25.0, 36.0, 49.0]
[8.0, 9.0] [64.0, 81.0]

Функции потерь и оптимизаторы

import torch
from torch import nn

logits = torch.tensor([[2.0, 0.5, -1.0]])
target = torch.tensor([0])
print(round(nn.CrossEntropyLoss()(logits, target).item(), 4))
pred = torch.tensor([2.5, 0.0])
true = torch.tensor([3.0, -0.5])
print(round(nn.MSELoss()(pred, true).item(), 4))

0.2413
0.25

Оптимизатор обновляет параметры по их градиентам. Самый простой — torch.optim.SGD: его метод step() применяет θ ← θ − η · g. Код ниже повторяет пример «одного шага» из первого урока (x = 2, y = 6, w = 1, η = 0,1) и получает тот же ответ w = 2,6. А zero_grad() сбрасывает градиенты в None.

import torch

w = torch.tensor([1.0], requires_grad=True)
optimizer = torch.optim.SGD([w], lr=0.1)
loss = ((2 * w - 6) ** 2).sum()
loss.backward()
print(w.grad)
optimizer.step()
print(w)
optimizer.zero_grad()
print(w.grad)

tensor([-16.])
tensor([2.6000], requires_grad=True)
None

Цикл обучения

Теперь настоящая задача. Сгенерируем точки в форме двух переплетённых полумесяцев («две луны»): по 500 точек в каждом классе, и прямой линией их не разделить — значит, нужна нелинейная модель. Делим данные на 800 обучающих и 200 тестовых примеров.

import math
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader, random_split

torch.manual_seed(42)

def make_moons(n, noise=0.2):
    t = torch.rand(n) * math.pi
    upper = torch.stack([torch.cos(t), torch.sin(t)], dim=1)
    lower = torch.stack([1 - torch.cos(t), 0.5 - torch.sin(t)], dim=1)
    X = torch.cat([upper, lower]) + noise * torch.randn(2 * n, 2)
    y = torch.cat([torch.zeros(n), torch.ones(n)]).long()
    return X, y

X, y = make_moons(500)
train_set, test_set = random_split(TensorDataset(X, y), [800, 200])
train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
test_loader = DataLoader(test_set, batch_size=100)
xb, yb = next(iter(train_loader))
print(len(train_set), len(test_set), len(train_loader))
print(xb.shape, yb.shape, yb.dtype)

800 200 25
torch.Size([32, 2]) torch.Size([32]) torch.int64

1. 1
   Обнули градиенты

   optimizer.zero_grad() — очисти градиенты предыдущего шага.

2. 2
   Прямой проход

   logits = model(xb) — предсказания для пакета.

3. 3
   Потери

   loss = loss_fn(logits, yb) — один скаляр.

4. 4
   Обратный проход

   loss.backward() — градиенты всех параметров.

5. 5
   Шаг

   optimizer.step() — обнови параметры.

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = nn.Sequential(
    nn.Linear(2, 16), nn.ReLU(),
    nn.Linear(16, 16), nn.ReLU(),
    nn.Linear(16, 2),
).to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for epoch in range(1, 21):
    model.train()
    total = 0.0
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        total += loss.item() * len(xb)
    if epoch == 1 or epoch % 5 == 0:
        print(f'epoch {epoch:2d}  loss {total / len(train_set):.4f}')

epoch  1  loss 0.4741
epoch  5  loss 0.1200
epoch 10  loss 0.0864
epoch 15  loss 0.0715
epoch 20  loss 0.0823

Потери быстро падают в первые эпохи, а потом колеблются около 0,07–0,09. Небольшой рост на 20-й эпохе — это нормально: мини-батчи случайны, поэтому шаги шумные. Складывая loss.item() * len(xb) и деля на N, мы получаем точное среднее потерь за эпоху (даже если последний пакет меньше).

Размер пакета B — тоже гиперпараметр. Большой пакет точнее оценивает градиент и лучше загружает GPU, но требует больше памяти и даёт меньше обновлений за эпоху. Маленький пакет шумный, однако этот шум иногда помогает обобщению. На практике начинают со значений от 32 до 256; при удвоении пакета обычно можно немного увеличить и скорость обучения.

В реальных проектах в цикл добавляют ещё два элемента. Планировщик скорости обучения (scheduler) со временем уменьшает η: например, torch.optim.lr_scheduler.CosineAnnealingLR, а scheduler.step() вызывают в конце каждой эпохи. Ранняя остановка (early stopping) после каждой эпохи измеряет потери на валидации, сохраняет лучшую модель и прекращает обучение, если потери не улучшаются несколько эпох подряд, — это самая простая защита от переобучения.

Оценка: режим eval и точность

При оценке важны две вещи. model.eval() меняет поведение некоторых слоёв: Dropout отключается, а BatchNorm использует статистики, накопленные при обучении. torch.no_grad() отключает построение графа. Это разные задачи, и обычно нужны обе сразу. Главная метрика в классификации — точность (accuracy):

def accuracy(model, loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            correct += (model(xb).argmax(dim=1) == yb).sum().item()
    return correct / len(loader.dataset)

print(f'train accuracy: {accuracy(model, train_loader):.3f}')
print(f'test accuracy:  {accuracy(model, test_loader):.3f}')

train accuracy: 0.974
test accuracy:  0.975

Сохранение и загрузка модели

model.state_dict() — словарь «имя параметра → тензор». Рекомендуется сохранять именно этот словарь, а не весь объект. Чтобы загрузить, заново создают модель с той же архитектурой и вызывают load_state_dict; map_location позволяет открыть на CPU файл, сохранённый на GPU.

torch.save(model.state_dict(), 'moons.pt')
restored = nn.Sequential(
    nn.Linear(2, 16), nn.ReLU(),
    nn.Linear(16, 16), nn.ReLU(),
    nn.Linear(16, 2),
).to(device)
restored.load_state_dict(torch.load('moons.pt', map_location=device))
print(list(restored.state_dict().keys()))
print(f'restored test accuracy: {accuracy(restored, test_loader):.3f}')

['0.weight', '0.bias', '2.weight', '2.bias', '4.weight', '4.bias']
restored test accuracy: 0.975

Чтобы прервать долгое обучение и потом продолжить, сохраняют чекпойнт: torch.save({'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict()}, 'ckpt.pt'). Статистики импульса Adam тоже хранятся в state_dict оптимизатора — без них продолжение было бы «холодным стартом».