PyTorch на практике: урок по Глубокое обучение

На собесе ML-инженера часто дают кусок кода обучения и просят найти ошибки. Почти все они из короткого списка: забытый zero_grad, eval или no_grad, softmax перед CrossEntropyLoss, накопление тензора loss, тензоры на разных устройствах.

Тензоры и autograd

Тензор это многомерный массив на CPU или GPU. Если у тензора requires_grad=True, операции над ним записываются в динамический граф. loss.backward() проходит граф назад и прибавляет градиенты к .grad параметров. Граф строится заново на каждом forward.

nn.Module

import torch
from torch import nn

class MLP(nn.Module):
    def __init__(self, d_in, d_hidden, n_classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, d_hidden), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(d_hidden, n_classes),
        )

    def forward(self, x):
        return self.net(x)  # логиты, без softmax

Модель вызывают как model(x). Параметры доступны через model.parameters(), сохраняют state_dict.

Цикл обучения и валидации

device = "cuda" if torch.cuda.is_available() else "cpu"
model = MLP(784, 256, 10).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for epoch in range(epochs):
    model.train()
    for x, y in train_loader:                 # shuffle=True
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        running += loss.item()

    model.eval()
    correct = 0
    with torch.no_grad():
        for x, y in val_loader:
            pred = model(x.to(device)).argmax(1)
            correct += (pred == y.to(device)).sum().item()
    acc = correct / len(val_loader.dataset)

Чек-лист багов

Ошибка	Последствие
Нет zero_grad	градиенты копятся между шагами
softmax перед CrossEntropyLoss	двойной softmax, слабое обучение
Нет model.eval()	dropout и BatchNorm в режиме обучения на валидации
Нет no_grad на валидации	лишняя память и время
total += loss без .item()	держим графы всех шагов, утечка памяти
shuffle=False на трейне	однородные батчи
Метки на CPU, модель на GPU	ошибка устройств
Деление на число батчей	неверная метрика при неполном последнем батче

model.eval() и torch.no_grad() это разные вещи: первое меняет поведение слоёв, второе отключает граф. torch.inference_mode() ещё строже для чистого инференса.

Заморозка слоёв

requires_grad = False у параметров бэкбона и оптимизатор только с параметрами головы. model.eval() веса не замораживает.

Данные

Dataset возвращает один объект, DataLoader собирает батчи.
num_workers > 0 и pin_memory=True, если GPU простаивает в ожидании данных.
Предобработка (нормализация, target encoding) обучается только на трейне.

Память GPU

На параметр при AdamW в fp32: вес 4 байта, градиент 4, два момента 8. Около 16 ГБ на миллиард параметров без активаций. При нехватке памяти:

меньше батч и накопление градиентов;
bf16 через torch.autocast;
gradient checkpointing;
поиск утечек (тензоры с графом в списках);
8-битные оптимизаторы, LoRA, FSDP.

Скорость

torch.compile захватывает граф и генерирует слитые ядра; первый запуск медленный из-за компиляции.
Эффективные ядра attention, mixed precision, отсутствие лишних синхронизаций (.item() в горячем цикле).

Воспроизводимость

Фиксируйте сиды (torch, numpy, random), сохраняйте конфиг, версии и чекпоинты: state_dict модели, оптимизатора, scheduler и номер эпохи.

Что спросят на собесе

Найдите ошибки в цикле обучения.
Разница model.eval() и torch.no_grad().
Почему растёт память во время обучения.
Как заморозить часть модели.
Сколько памяти нужно для обучения модели на N параметров.

PyTorch на практике за 8 минут