nn.Module, оптимизаторы, цикл обучения

Зачем это нужно

В прошлом уроке мы обучили линейную регрессию с двумя параметрами: сами создали w и b, сами обновили их внутри torch.no_grad(), сами обнулили градиенты. Для двух параметров это нормально. В MiniMind их десятки миллионов, в DeepSeek-V3 — 671 миллиард, разложенные по тысячам тензоров. Перечислять их вручную невозможно.

PyTorch решает это тремя кирпичами:

• nn.Module — контейнер, который сам находит все параметры модели, сколько бы слоёв ни было вложено;
• torch.optim — оптимизаторы из урока 1.4 (SGD, Adam, AdamW), уже написанные и проверенные;
• Dataset и DataLoader — подача данных батчами.

Из них собирается цикл обученияцикл обученияEN training loopПовторять: прямой проход и loss → обнуление градиентов → backward → шаг оптимизатора. — пять строк, которые одинаковы для классификатора цифр и для LLM. В конце урока вы обучите классификатор рукописных цифр MNISTMNIST (рукописные цифры)EN MNIST70 000 изображений рукописных цифр 28×28 (60 000 обучающих и 10 000 тестовых). «Hello, world» глубокого обучения. и сохраните его в чекпоинт.

nn.Module: модель как дерево слоёв

В уроке 0.2.3 вы писали мини-версию nn.Module: родительский класс с __call__, который вызывает forward. Настоящий nn.Module делает то же плюс главное — регистрирует параметры. Любой слой, присвоенный атрибуту self.что-то, становится частью модели:

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self, n_in=784, hidden=128, n_out=10):
        super().__init__()                       # обязательно первой строкой!
        self.fc1 = nn.Linear(n_in, hidden)       # слой-атрибут → его параметры видны модели
        self.fc2 = nn.Linear(hidden, n_out)

    def forward(self, x):                        # описываем только прямой проход
        h = torch.relu(self.fc1(x))              # backward сделает autograd
        return self.fc2(h)                       # логиты

torch.manual_seed(0)
model = TinyMLP()
print(model)
for name, p in model.named_parameters():
    print(f"{name:12} {tuple(p.shape)}")
print("всего параметров:", sum(p.numel() for p in model.parameters()))

x = torch.randn(32, 784)                         # батч из 32 «картинок»
print(model(x).shape)                            # вызываем model(x), а не model.forward(x)

TinyMLP(
  (fc1): Linear(in_features=784, out_features=128, bias=True)
  (fc2): Linear(in_features=128, out_features=10, bias=True)
)
fc1.weight   (128, 784)
fc1.bias     (128,)
fc2.weight   (10, 128)
fc2.bias     (10,)
всего параметров: 101770
torch.Size([32, 10])

Проверим число параметров руками: $784 \cdot 128 + 128 + 128 \cdot 10 + 10 = 100\,352 + 128 + 1\,280 + 10 = 101\,770$. Почти все параметры — в первом слое. Это типично: самые «толстые» матрицы стоят там, где широкий вход.

Что даёт модульмодуль (nn.Module)EN module (nn.Module)Базовый класс всех слоёв и моделей PyTorch. Сам находит параметры во вложенных модулях, переносит их на устройство, сохраняет и загружает. В d2l-zh называется «блоком».:

Метод	Что делает
model.parameters()	все обучаемые тензоры, включая вложенные модули, — их отдают оптимизатору
model.named_parameters()	то же с именами: fc1.weight, fc2.bias…
model.to(device)	переносит все параметры на CPU или GPU; у модели работает на месте
model.train() / model.eval()	режим обучения или оценки — важно для dropout и BatchNorm (урок 3.4)
model.state_dict()	словарь «имя → тензор» для сохранения на диск
model(x)	вызывает forward плюс служебные «хуки» (функции, которые PyTorch запускает до и после forward); forward напрямую не вызывают

Модули вкладываются друг в друга, как матрёшки. TinyMLP содержит два nn.Linear, блок трансформера содержит attention и MLP, а вся LLM — embedding, список блоков и выходной слой. parameters() обходит всё дерево.

nn.Linear изнутри

nn.Linear(in, out) хранит матрицу weight формы (out, in) и вектор bias формы (out,). Считает он $y = x W^\top + b$:

lin = nn.Linear(3, 2)                 # 3 входа → 2 выхода
print(lin.weight.shape, lin.bias.shape)          # torch.Size([2, 3]) torch.Size([2])
x = torch.randn(4, 3)                            # батч из 4 примеров
print(torch.allclose(lin(x), x @ lin.weight.T + lin.bias))   # True

Порядок «(выходы, входы)» путает почти всех, особенно при чтении чужих чекпоинтов. Веса инициализируются случайно — небольшими числами порядка $1/\sqrt{\text{in}}$. Почему именно так и что будет при плохой инициализации — в уроке «Инициализация и нормализация».

Две ловушки регистрации

nn.Module находит только то, что «знает». Два классических бага:

class Buggy(nn.Module):
    def __init__(self):
        super().__init__()
        self.scale = torch.ones(1, requires_grad=True)          # обычный тензор — НЕ параметр
        self.layers = [nn.Linear(4, 4) for _ in range(3)]       # обычный список — НЕ виден

class Fixed(nn.Module):
    def __init__(self):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(1))                 # nn.Parameter регистрируется
        self.layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)])

count = lambda m: sum(p.numel() for p in m.parameters())
print("Buggy — параметров:", count(Buggy()))   # Buggy — параметров: 0
print("Fixed — параметров:", count(Fixed()))   # Fixed — параметров: 61

У Buggy ноль параметров. Оптимизатор не будет его учить, .to("cuda") его не перенесёт, state_dict() не сохранит. При этом никакой ошибки нет, и модель даже что-то выдаёт. Правила:

• свой обучаемый тензор → nn.Parameter(...);
• список слоёв → nn.ModuleList, словарь слоёв → nn.ModuleDict;
• super().__init__() — первой строкой, иначе PyTorch упадёт при первом же присваивании слоя.

После описания модели всегда печатайте число параметров и сверяйте с расчётом на бумаге.

nn.Sequential: когда слои идут цепочкой

Если forward просто пропускает данные через слои по очереди, класс можно не писать:

model = nn.Sequential(
    nn.Flatten(),            # (B, 28, 28) → (B, 784)
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)
print(model[1].weight.shape)  # к слоям обращаются по индексу: torch.Size([128, 784])

Для LLM Sequential мало: там есть остаточные связи и несколько входов. Поэтому реальные модели пишут классами, а Sequential используют для простых кусков внутри, например для MLP в блоке трансформера.

Функция потерь: на вход — логиты

Для классификации берут cross-entropy из урока 1.3:

import torch.nn.functional as F

loss = F.cross_entropy(logits, y)    # logits: (B, C) float, y: (B,) int64 — номера классов

Три правила, которые нужно запомнить:

1. В F.cross_entropy передают логиты. Softmax и логарифм она применит сама, устойчиво, через log-sum-exp (урок 1.3.2). Модель поэтому заканчивается nn.Linear, без softmax.
2. Метки — целые числа типа int64 (torch.long), а не one-hot векторы. Для цифры «7» метка — просто 7.
3. По умолчанию loss усредняется по батчу (reduction="mean"). Поэтому backward() вызывают от одного числа.

Есть и классовый вариант nn.CrossEntropyLoss(): создаёте объект один раз и вызываете как функцию. Разницы в вычислениях нет.

Оптимизаторы: ваш Adam из урока 1.4 — в одну строку

В уроке 1.4 вы писали SGD, momentum и Adam руками. torch.optim содержит те же формулы. Оптимизатор получает список параметров и делает две вещи: zero_grad() обнуляет их .grad, step() обновляет по формуле своего алгоритма. Проверим, что SGD делает ровно $w \leftarrow w - \eta \nabla w$:

torch.manual_seed(0)
lin = nn.Linear(2, 1)
opt = torch.optim.SGD(lin.parameters(), lr=0.1)

x, y = torch.tensor([[1.0, 2.0]]), torch.tensor([[1.0]])
loss = ((lin(x) - y) ** 2).mean()
opt.zero_grad()
loss.backward()

expected = lin.weight - 0.1 * lin.weight.grad    # что сделал бы наш ручной SGD
opt.step()                                       # что делает PyTorch
print(torch.allclose(lin.weight, expected))      # True

Оптимизатор	Строка	Типичный lr	Где встречается
SGD	optim.SGD(params, lr=0.1)	0.01–0.1	учебные примеры, иногда CV
SGD + momentum	optim.SGD(params, lr=0.1, momentum=0.9)	0.01–0.1	ResNet и классическое CV
Adam	optim.Adam(params, lr=1e-3)	$10^{-4}$–$10^{-3}$	«по умолчанию» для небольших сетей
AdamW	optim.AdamW(params, lr=1e-3, weight_decay=0.01)	$10^{-4}$–$10^{-3}$	стандарт для трансформеров и LLM

Диапазоны lr — ориентиры, а не законы. Для больших LLM learning rate обычно ещё меньше, порядка $10^{-4}$ и ниже. Текущий lr лежит в opt.param_groups[0]["lr"]. Через param_groups меняют lr по расписанию (warmup и cosine из урока 1.4) — так делает MiniMind.

Данные: Dataset и DataLoader

Dataset — это то, что вы уже писали в уроке 0.2.3: класс с двумя методами. __len__ возвращает число примеров, __getitem__(i) — i-й пример (обычно пару «вход, метка»). Для данных, которые уже лежат в тензорах, есть готовый TensorDataset.

загрузчик данныхзагрузчик данных (DataLoader)EN data loaderБерёт Dataset и выдаёт батчи: перемешивает, склеивает примеры в тензоры, может грузить данные в нескольких процессах. берёт Dataset и выдаёт батчи: перемешивает индексы, достаёт примеры, склеивает их в один тензор с новой осью батча.

from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)
X = torch.randn(1000, 20)                 # 1000 примеров по 20 признаков
y = torch.randint(0, 3, (1000,))          # метки трёх классов
ds = TensorDataset(X, y)                  # готовый Dataset поверх тензоров
print(len(ds), ds[0][0].shape, ds[0][1])  # 1000 torch.Size([20]) tensor(2) — ds[0] это пара (вход, метка)

loader = DataLoader(ds, batch_size=64, shuffle=True)
print("батчей в эпохе:", len(loader))     # 16
xb, yb = next(iter(loader))               # взять только первый батч (iter/next — «дай следующий элемент»)
print(xb.shape, yb.shape)                 # torch.Size([64, 20]) torch.Size([64])
print([len(b) for _, b in loader][-3:])   # размеры трёх последних батчей: [64, 64, 40] — последний неполный

Словарь, который понадобится дальше:

• батч (мини-батчмини-батчEN mini-batchНебольшая случайная порция примеров, по которой считается одна оценка градиента и делается один шаг.) — порция примеров на один шаг оптимизатора; размер батчаразмер батчаEN batch sizeЧисло примеров (у LLM — часто число токенов) на один шаг оптимизатора. Больше батч — меньше шум градиента, но дороже шаг. — сколько их в порции;
• шаг (step, iteration) — один вызов opt.step();
• эпохаэпохаEN epochОдин полный проход по обучающим данным. LLM на претрейне часто не делают и одной эпохи. — полный проход по всем данным. Шагов в эпохе: $\lceil N / B \rceil$, где скобки $\lceil \cdot \rceil$ означают округление вверх до целого. Для 1000 примеров и батча 64 это $\lceil 15.6 \rceil = 16$: пятнадцать полных батчей и один из 40 примеров.

Полезные параметры DataLoader: shuffle=True для обучения (на валидации и тесте не нужно), drop_last=True — выбросить неполный последний батч, num_workers=4 — готовить батчи в отдельных процессах, пока модель считает, pin_memory=True — ускоряет копирование на GPU.

Стандартный цикл обучения

Теперь всё собирается в один цикл. Запомните его наизусть — вы будете писать его десятки раз:

flowchart LR
D["DataLoader<br/>батч x, y"] --> F["① logits = model(x)"]
F --> L["② loss = F.cross_entropy(logits, y)"]
L --> Z["③ opt.zero_grad()"]
Z --> B["④ loss.backward()"]
B --> S["⑤ opt.step()"]
S -->|"следующий батч"| D
S -.->|"раз в эпоху"| V["оценка на валидации<br/>model.eval() + no_grad"]

Цикл обучения: пять строк на каждый батч и оценка раз в эпоху

for epoch in range(n_epochs):
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        logits = model(x)                    # 1. прямой проход
        loss = F.cross_entropy(logits, y)    # 2. функция потерь
        opt.zero_grad()                      # 3. стереть градиенты прошлого шага
        loss.backward()                      # 4. посчитать новые
        opt.step()                           # 5. сдвинуть веса
    val_loss, val_acc = evaluate(model, val_loader, device)

Порядок zero_grad → backward → step важен. zero_grad() можно вызвать и в начале итерации — главное, чтобы между двумя backward() градиенты обнулялись. По умолчанию zero_grad() не заполняет .grad нулями, а ставит None: так экономится память и время.

Оценку пишут отдельной функцией с двумя «переключателями». Строка @torch.no_grad() над def — декоратор из урока 0.2.3: он включает режим «без графа» на всё время работы функции:

@torch.no_grad()                  # 1) не строить граф: меньше памяти, быстрее
def evaluate(model, loader, device):
    model.eval()                  # 2) режим оценки для dropout/BatchNorm
    ...
    model.train()                 # вернуть режим обучения

Это две разные вещи, и нужны обе. no_grad отключает autograd, а eval() меняет поведение отдельных слоёв. Dropout с p=0.5 в режиме train случайно обнуляет в среднем половину входов, а в режиме eval пропускает всё как есть.

Чекпоинты: сохраняем и продолжаем

Обучение LLM длится дни и недели, и сбои неизбежны: падает узел, кончается место на диске. чекпоинтчекпоинт (контрольная точка)EN checkpointСохранённое состояние обучения: веса модели, состояние оптимизатора, номер шага. Позволяет продолжить обучение после сбоя. — снимок, из которого обучение продолжается с того же места. Его основа — словарь состояниясловарь состояния (state_dict)EN state dictСловарь «имя → тензор» со всеми параметрами и буферами модуля (или с состоянием оптимизатора). Именно его сохраняют через torch.save.:

ckpt = {
    "model": model.state_dict(),        # веса: {"net.1.weight": тензор, ...}
    "optimizer": opt.state_dict(),      # моменты Adam для каждого параметра
    "epoch": epoch,                     # где остановились
}
torch.save(ckpt, "mnist_mlp.pt")

# ...позже, в другом процессе:
ckpt = torch.load("mnist_mlp.pt")       # weights_only=True по умолчанию (PyTorch ≥ 2.6)
model = MLP()                           # сначала такая же архитектура (класс MLP — в скрипте ниже)...
model.load_state_dict(ckpt["model"])    # ...потом веса
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
opt.load_state_dict(ckpt["optimizer"])
start_epoch = ckpt["epoch"] + 1

Сохраняют state_dict, а не весь объект модели. Словарь тензоров не зависит от вашего кода: его можно загрузить в исправленную версию класса или в другой проект. Поэтому с версии 2.6 torch.load по умолчанию читает только тензоры и простые типы (weights_only=True). Внутри torch.save использует pickle — стандартный формат Python для сохранения любых объектов. Беда в том, что при загрузке pickle-файл может выполнить произвольный код: скачали чужой чекпоинт — запустили чужую программу. Файл с одними тензорами такой опасности не несёт. Открытые веса Qwen и DeepSeek публикуют в формате safetensors: в нём вообще нет исполняемого кода.

Зачем состояние оптимизатора? AdamW хранит на каждый параметр два тензора той же формы: exp_avg (первый момент — скользящее среднее градиентов) и exp_avg_sq (второй — скользящее среднее их квадратов); вы писали их руками в уроке 1.4.2. Загрузите только веса — Adam начнёт с нулевых моментов, и первые шаги будут резкими. Отсюда полезное правило: чекпоинт для продолжения обучения с AdamW в 3 раза больше самих весов (в одинаковом dtype). Наш MLP: 203 530 параметров × 4 байта ≈ 0.8 МБ весов, а файл чекпоинта — 2.4 МБ. Для модели на 7 млрд параметров в fp32 это 28 ГБ весов и 84 ГБ чекпоинта.

Практика: MNIST на CPU

MNIST — 70 000 рукописных цифр размером 28×28 пикселей: 60 000 для обучения и 10 000 для теста. Соберём всё в один скрипт mnist_mlp.py. Данные скачаем сами (4 файла, около 11 МБ) и напишем свой Dataset.

Файлы MNIST лежат в простом бинарном формате IDX. В начале файла заголовок: у картинок 16 байт («магическое число» — метка формата, количество картинок, 28, 28), у меток 8 байт. Дальше идут сырые байты 0–255: по одному байту на пиксель (яркость от 0 до 255, урок 0.0.2). Поэтому весь «парсер» — одна строка np.frombuffer(..., offset=16): «прочитай байты как массив чисел, пропустив первые 16».

Скрипт длинный, но устроен просто: четыре блока, отмеченные комментариями-разделителями, — данные, модель, оценка, обучение. Читайте его по блокам сверху вниз; каждый блок вы уже видели в этом уроке по отдельности.

"""MNIST-классификатор на PyTorch: данные → модель → цикл обучения → чекпоинт."""
import gzip
import time
import urllib.request
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, random_split

# ── 1. Данные ────────────────────────────────────────────────────────────────
URL = "https://ossci-datasets.s3.amazonaws.com/mnist/"  # зеркало, которым пользуется torchvision
FILES = {
    "train": ("train-images-idx3-ubyte.gz", "train-labels-idx1-ubyte.gz"),
    "test": ("t10k-images-idx3-ubyte.gz", "t10k-labels-idx1-ubyte.gz"),
}
DATA = Path("data/mnist")


def download():
    """Скачивает 4 файла MNIST (≈ 11 МБ), если их ещё нет."""
    DATA.mkdir(parents=True, exist_ok=True)
    for pair in FILES.values():
        for name in pair:
            if not (DATA / name).exists():
                print("скачиваю", name)
                urllib.request.urlretrieve(URL + name, DATA / name)


def read_idx(path, header_bytes):
    """Формат IDX: заголовок фиксированной длины, дальше сырые байты 0..255."""
    with gzip.open(path, "rb") as f:
        return np.frombuffer(f.read(), dtype=np.uint8, offset=header_bytes)


class MNIST(Dataset):
    """Dataset = «умеет сказать длину и выдать пример по индексу»."""

    MEAN, STD = 0.1307, 0.3081  # среднее и std пикселей обучающей выборки

    def __init__(self, split="train"):
        img_file, lbl_file = FILES[split]
        images = read_idx(DATA / img_file, 16).reshape(-1, 28, 28)
        labels = read_idx(DATA / lbl_file, 8)
        x = torch.from_numpy(images.copy()).float() / 255.0  # uint8 → float32 в [0, 1]
        self.x = (x - self.MEAN) / self.STD                  # стандартизация
        self.y = torch.from_numpy(labels.copy()).long()      # метки int64 — так хочет cross_entropy

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

    def __getitem__(self, i):
        return self.x[i], self.y[i]  # картинка (28, 28) и метка-число


# ── 2. Модель ────────────────────────────────────────────────────────────────
class MLP(nn.Module):
    def __init__(self, hidden=256):
        super().__init__()                # обязательно первой строкой!
        self.net = nn.Sequential(
            nn.Flatten(),                 # (B, 28, 28) → (B, 784)
            nn.Linear(784, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 10),        # 10 логитов — по одному на цифру
        )

    def forward(self, x):
        return self.net(x)                # логиты, НЕ вероятности


# ── 3. Оценка ────────────────────────────────────────────────────────────────
@torch.no_grad()                          # граф не нужен — экономим память и время
def evaluate(model, loader, device):
    model.eval()                          # режим оценки (важно для dropout/BatchNorm)
    loss_sum, correct, total = 0.0, 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        logits = model(x)
        loss_sum += F.cross_entropy(logits, y, reduction="sum").item()
        correct += (logits.argmax(dim=1) == y).sum().item()
        total += len(y)
    model.train()                         # вернуть режим обучения
    return loss_sum / total, correct / total


# ── 4. Обучение ──────────────────────────────────────────────────────────────
def main(epochs=3, batch_size=128, lr=1e-3):
    torch.manual_seed(0)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    download()

    full = MNIST("train")
    train_ds, val_ds = random_split(full, [55_000, 5_000],
                                    generator=torch.Generator().manual_seed(0))
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=1000)
    test_loader = DataLoader(MNIST("test"), batch_size=1000)

    model = MLP().to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"устройство: {device}, параметров: {n_params:,}, батчей в эпохе: {len(train_loader)}")

    for epoch in range(1, epochs + 1):
        t0 = time.perf_counter()
        for step, (x, y) in enumerate(train_loader):
            x, y = x.to(device), y.to(device)
            logits = model(x)                    # 1. прямой проход
            loss = F.cross_entropy(logits, y)    # 2. функция потерь
            opt.zero_grad()                      # 3. стереть старые градиенты
            loss.backward()                      # 4. обратный проход
            opt.step()                           # 5. шаг оптимизатора
            if step % 200 == 0:
                print(f"  эпоха {epoch}, шаг {step:3d}: loss {loss.item():.4f}")
        val_loss, val_acc = evaluate(model, val_loader, device)
        print(f"эпоха {epoch}: val loss {val_loss:.4f}, val accuracy {val_acc:.4f}, "
              f"{time.perf_counter() - t0:.1f} с")

    # Чекпоинт: веса + состояние оптимизатора + номер эпохи
    torch.save({"model": model.state_dict(),
                "optimizer": opt.state_dict(),
                "epoch": epochs}, "mnist_mlp.pt")
    test_loss, test_acc = evaluate(model, test_loader, device)
    print(f"test accuracy {test_acc:.4f} — чекпоинт сохранён в mnist_mlp.pt")


if __name__ == "__main__":   # при import этот блок не выполнится — пригодится в следующем уроке
    main()

Запуск: python mnist_mlp.py. Ожидаемый вывод (время у вас будет своё; последние знаки loss могут отличаться на другой версии PyTorch или другом процессоре):

устройство: cpu, параметров: 203,530, батчей в эпохе: 430
  эпоха 1, шаг   0: loss 2.3326
  эпоха 1, шаг 200: loss 0.1087
  эпоха 1, шаг 400: loss 0.1463
эпоха 1: val loss 0.1618, val accuracy 0.9494, … с
  эпоха 2, шаг   0: loss 0.1820
  …
эпоха 2: val loss 0.1092, val accuracy 0.9666, … с
  …
эпоха 3: val loss 0.1018, val accuracy 0.9700, … с
test accuracy 0.9721 — чекпоинт сохранён в mnist_mlp.pt

Три эпохи — и 97.2% на тесте, то есть 279 ошибок из 10 000. Обратите внимание на первую строку: loss 2.33 на нулевом шаге. Это почти ровно $\ln 10 \approx 2.30$ — loss модели, которая пока угадывает равномерно из 10 классов (урок 1.3). В следующем уроке это станет первой проверкой здравого смысла.

Несколько деталей кода, которые легко пропустить:

• images.copy() нужен потому, что np.frombuffer даёт массив только для чтения, а PyTorch не любит делить память с таким массивом (вспомните from_numpy из прошлого урока).
• Стандартизация (x - 0.1307) / 0.3081 делает входы со средним 0 и разбросом 1, как в уроке 2.2. Без неё обучение тоже идёт, но медленнее и капризнее к lr.
• random_split с собственным генератором отрезает 5 000 примеров на валидацию — всегда одних и тех же. Тест мы трогаем один раз, в самом конце (урок 2.1).
• f"{n_params:,}" печатает число с запятыми между тысячами: 203,530.
• weight_decay=0.01 в AdamW — «затухание весов» из урока 1.4.2; подробнее о нём — в уроке 3.4.2.
• Сколько ждать эпоху? Прикинем, как в прошлом уроке. Прямой проход одного батча — это $128 \cdot (784 \cdot 256 + 256 \cdot 10) \approx 26$ млн умножений со сложением, обратный — примерно вдвое больше. За 430 шагов набирается около $7 \cdot 10^{10}$ операций, и процессор с десятками гигафлопс справится за секунды. Если эпоха идёт минутами, скорее всего процессор занят другими программами и потоки PyTorch мешают друг другу. Добавьте в начало main() строку torch.set_num_threads(2) или даже (1). У нас на четырёхъядерной машине, загруженной другими задачами, переход с 4 потоков (по умолчанию) на 1 ускорил шаг обучения в сотни раз.

Альтернатива: torchvision

В туториалах MNIST обычно грузят одной строкой из пакета torchvision (pip install torchvision, для CPU — с тем же --index-url, что и torch): torchvision.datasets.MNIST("data", train=True, download=True, transform=transforms.ToTensor()). Мы написали загрузчик сами, чтобы было видно, что внутри Dataset нет никакой магии. И чтобы курс не требовал лишней зависимости.

Задания:

1. Тетрадь ошибок для модели. Найдите все ошибки модели на тесте и нарисуйте первые 16 (код ниже). Посмотрите на них глазами. Какие цифры путаются чаще? Сколько из них вы бы сами прочитали уверенно?
2. Продолжить обучение. Напишите resume.py: загрузите чекпоинт, восстановите модель и оптимизатор, обучите ещё 2 эпохи и сохраните новый чекпоинт с epoch = 5. Сравните с запуском main(epochs=5) с нуля. Совпали ли числа до последнего знака? Если нет, подумайте, что ещё случайно в обучении, кроме начальных весов. Чего не хватает в чекпоинте, разберём в следующем уроке, а в ДЗ модуля такое совпадение проверяет тест.
3. Эксперимент. Поменяйте hidden=256 на 64 и на 1024 и lr на 1e-2 и 1e-4. Сведите результаты в таблицу «параметров — val accuracy — время эпохи». Меняйте одну вещь за раз.

import matplotlib.pyplot as plt
import torch
from mnist_mlp import MLP, MNIST          # импортируем из своего скрипта

model = MLP()
model.load_state_dict(torch.load("mnist_mlp.pt")["model"])
model.eval()

test = MNIST("test")
with torch.no_grad():
    pred = model(test.x).argmax(dim=1)    # все 10 000 картинок одним батчем
wrong = (pred != test.y).nonzero().squeeze(1)
print("ошибок:", len(wrong), "из", len(test))   # ошибок: 279 из 10000

fig, axes = plt.subplots(2, 8, figsize=(12, 3.5))
for ax, i in zip(axes.flat, wrong[:16].tolist()):
    img = test.x[i] * MNIST.STD + MNIST.MEAN    # обратно в [0, 1] для картинки
    ax.imshow(img, cmap="gray")
    ax.set_title(f"верно {test.y[i].item()}, модель {pred[i].item()}", fontsize=8)
    ax.axis("off")
plt.tight_layout()
plt.savefig("mistakes.png")

У нас среди первых 16 ошибок больше всего девяток, которые модель приняла за 8 или 4. Часть цифр написана так, что и человек засомневается. Свёрточные сети (урок 3.4) на MNIST обычно доходят до 99% и выше: они «знают», что соседние пиксели связаны, а наш MLP видит картинку как плоский список из 784 чисел.

🖥 На вашей машине (4×GPU)

Скрипт сам выберет cuda, если видеокарта доступна. Для одной карты запускайте CUDA_VISIBLE_DEVICES=0 python mnist_mlp.py. На таком маленьком MLP GPU ускорит обучение ненамного. Сам шаг на видеокарте занимает доли миллисекунды, и время уходит на подготовку батчей на CPU и запуск десятков мелких операций. Попробуйте увеличить batch_size до 1024 (и lr до 3e-3), добавить в DataLoader num_workers=8, pin_memory=True — и сравните время эпохи на CPU и GPU. Это первый опыт «профилирования»: искать, где на самом деле тратится время. Настоящий профайлер и устройство GPU — в уроке 5.7.1.

Китайский угол: MiniMind и DeepSeek-V3 — те же кирпичи

MiniMind (github.com/jingyaogong/minimind, Apache-2.0) — открытый китайский проект, где LLM обучают с нуля на одной видеокарте. Откройте trainer/train_pretrain.py — вы узнаете каждую строку этого урока:

• оптимизатор: optim.AdamW(model.parameters(), lr=args.learning_rate);
• данные: свой PretrainDataset и DataLoader(train_ds, ..., num_workers=args.num_workers, pin_memory=True);
• шаг: backward() → clip_grad_norm_ → optimizer.step() → optimizer.zero_grad(set_to_none=True). В оригинале вызовы обёрнуты в scaler для смешанной точности (урок 5.4), а шаг оптимизатора делается раз в несколько батчей;
• learning rate меняется по расписанию через optimizer.param_groups;
• чекпоинт — state_dict модели, переведённый в fp16 на CPU, плюс отдельный файл для продолжения обучения с optimizer.state_dict(), epoch и step. Пишется сначала во временный .tmp, затем os.replace: если процесс упадёт посреди записи, старый чекпоинт не испортится.

Подсказки к аргументам там на китайском: число эпох буквально названо «числом раундов обучения», а остальное переведено дословно — начальный learning rate, шаги накопления градиента, порог обрезки градиента, случайное зерно (seed). Машинный перевод с ними справляется. По умолчанию batch_size=32, learning_rate=5e-4, grad_clip=1.0. Клиппинг и накопление градиента разберём в уроке «Цикл обучения LLM».

DeepSeek-V3. Открытый код инференса модели на 671B параметров — это классы MLA(nn.Module), MoE(nn.Module), Gate(nn.Module), Block(nn.Module) и Transformer(nn.Module), где блоки лежат в torch.nn.ModuleList. RMSNorm хранит свой вес как nn.Parameter(torch.ones(dim)) — ровно правило из раздела про ловушки регистрации. Устроено так же, как наш TinyMLP, только дерево модулей несравнимо больше.

Учебники. В d2l-zh модули называются «блоками»: разделы 5.1 «Слои и блоки», 5.2 «Управление параметрами» и 5.5 «Чтение и запись файлов» покрывают этот урок, с кодом на PyTorch и PaddlePaddle. В «PyTorch: глубоко и просто» (thorough-pytorch) от Datawhale глава 3 разбирает те же компоненты (чтение данных, построение модели, функция потерь, обучение и оценка, оптимизатор), а в главе 4 «Практика по основам PyTorch» есть полный пример 4.4 «Классификация изображений FashionMNIST». Fashion-MNIST — «модный» двойник MNIST: те же 28×28 пикселей и 10 классов, только вместо цифр одежда и обувь.

Проверь себя

Проверь себя6 вопросов·

1 Сколько обучаемых параметров у nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))?

2 В __init__ модели написано self.blocks = [Block() for _ in range(4)]. Что произойдёт при обучении?

PyTorch выдаст ошибку при создании моделиВсё будет работать как с nn.ModuleListОшибки не будет, но веса блоков не попадут в parameters(): оптимизатор их не обучит, .to() не перенесёт, state_dict() не сохранитБлоки обучатся, но не сохранятся в чекпоинт

3 Какие утверждения о цикле обучения верны?Можно выбрать несколько ответов

В F.cross_entropy передают логиты, а не результат softmaxtorch.no_grad() при оценке переключает dropout в режим инференсаМетки классов для F.cross_entropy — целые числа типа int64Чекпоинт с состоянием AdamW примерно в 3 раза больше самих весов

4 Обучающая выборка — 50 000 примеров, batch_size=64, drop_last=False. Сколько раз за эпоху вызовется opt.step()?

5 Зачем сохранять в чекпоинт opt.state_dict(), если веса модели уже сохранены?

Без него нельзя загрузить весаВ нём моменты Adam (exp_avg, exp_avg_sq); без них оптимизатор после перезапуска начнёт с нулей и первые шаги будут резкимиВ нём хранится архитектура моделиЭто нужно только для SGD

6 Какая форма у nn.Linear(256, 10).weight?

(256, 10)(10, 256)(10,)(256,)

Проверить ответы

Итоги

• nn.Module регистрирует слои и nn.Parameter, присвоенные атрибутам. Списки — только через nn.ModuleList. После описания модели сверяйте число параметров с расчётом.
• torch.optim — те же SGD/Adam/AdamW, что вы писали руками; zero_grad() и step() работают со всеми параметрами сразу.
• Dataset — это __len__ + __getitem__; DataLoader режет данные на батчи и перемешивает их.
• Цикл обучения: model(x) → loss → zero_grad → backward → step. Оценка — в model.eval() и torch.no_grad().
• Чекпоинт = state_dict модели + state_dict оптимизатора + шаг. С AdamW он втрое больше весов.

Дальше: что делать, когда loss не падает. Переобучение одного батча, мониторинг градиентов, галерея типичных багов и китайская «алхимия» («ляньдань»).