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

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

Функция ошибки

Для классификации — nn.CrossEntropyLoss:

loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, labels)       # logits: (N, C), labels: (N,) номера классов

Она сама превращает логиты в вероятности (log-softmax) и сравнивает с правильными классами. Поэтому подавайте ей логиты, без softmax. Для предсказания чисел (регрессии) — nn.MSELoss.

Оптимизатор

Вместо ручного w -= lr * w.grad берут готовый оптимизатор:

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
# или чаще: torch.optim.Adam(model.parameters(), lr=1e-3)

Он знает все параметры модели и умеет обнулять их градиенты (zero_grad) и делать шаг (step). Adam сам подстраивает размер шага для каждого параметра и обычно учится быстрее.

Шаг обучения

for epoch in range(epochs):
    optimizer.zero_grad()            # 1. обнулить старые градиенты
    logits = model(x)                # 2. прямой проход
    loss = loss_fn(logits, y)        #    ошибка
    loss.backward()                  # 3. производные
    optimizer.step()                 # 4. шаг

Порядок важен: обнуление после backward() стирает только что посчитанные градиенты, а step() до backward() шагает по старым.

Режимы train и eval

Некоторые слои ведут себя по-разному при обучении и проверке. Например, nn.Dropout(0.5) при обучении случайно обнуляет половину чисел — это мешает сети «зазубрить» примеры, — а при проверке пропускает всё.

model.eval()                     # режим проверки
with torch.no_grad():            # без записи для производных — быстрее
    acc = (model(x_test).argmax(dim=1) == y_test).float().mean()
model.train()                    # обратно в режим обучения

model.eval() и torch.no_grad() — разные вещи: первое меняет поведение слоёв, второе выключает autograd. При проверке нужны оба.

Повторяемость

Начальные веса и перемешивание данных случайные. Чтобы результат повторялся, в начале программы задают torch.manual_seed(0).

Итог

nn.CrossEntropyLoss — на логитах и номерах классов;
torch.optim.SGD / Adam;
zero_grad → прямой проход и ошибка → backward → step;
проверка: model.eval() + torch.no_grad(), потом model.train().