PyTorch: Learn the Basics, часть 4 — цикл обучения и сохранение модели

Заключительная часть перевода курса Learn the Basics (исходники — pytorch/tutorials, BSD-3-Clause): разделы «Optimizing Model Parameters» и «Save and Load the Model». Предыдущая часть — модель и autograd, начало курса — быстрый старт и тензоры.

Как читать вместе с курсом ChinaAI

Цикл «предсказание → потери → backward → step → zero_grad» — центральная тема урока «nn.Module, оптимизаторы, цикл обучения». Что такое градиентный спуск и шаг обучения, объясняет урок «Градиентный спуск», а что делать, когда потери не падают, — урок «Как отлаживать обучение».

6. Оптимизация параметров модели

Теперь, когда у нас есть модель и данные, пора обучить, проверить и протестировать модель, оптимизируя её параметры на наших данных. Обучение модели — итеративный процесс. На каждой итерации модель делает догадку о выходе, вычисляет ошибку своей догадки (потери, loss), собирает производные ошибки по своим параметрам (как мы видели в предыдущем разделе) и оптимизирует эти параметры градиентным спуском. Более подробный разбор этого процесса — в видео 3Blue1Brown об обратном распространении.

Код из предыдущих разделов

Загружаем код из разделов Dataset и DataLoader и Сборка модели.

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import v2

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
)

train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork()

Гиперпараметры

ГиперпараметрыгиперпараметрEN hyperparameterНастройка, которую не учат градиентом, а выбирают по валидации — степень полинома, λ, learning rate, число слоёв. — настраиваемые параметры, которые управляют процессом оптимизации модели. Разные значения гиперпараметров могут влиять на обучение модели и скорость сходимости (подробнее о подборе гиперпараметров).

Для обучения зададим такие гиперпараметры:

• Число эпох — сколько раз пройти по всему набору данных.
• Размер батча — сколько примеров данных пропустить через сеть, прежде чем обновить параметры.
• Шаг обучения (learning rate) — насколько сильно обновлять параметры модели на каждом батче или эпохе. Маленькие значения дают медленное обучение, а большие могут привести к непредсказуемому поведению во время обучения.

learning_rate = 1e-3
batch_size = 64
epochs = 5

Цикл оптимизации

Задав гиперпараметры, мы можем обучать и оптимизировать модель в цикле оптимизации. Каждая итерация цикла оптимизации называется эпохой.

Каждая эпоха состоит из двух главных частей:

• Цикл обучения — проходим по обучающему набору и пытаемся приблизиться к оптимальным параметрам.
• Цикл проверки/тестирования — проходим по тестовому набору и проверяем, улучшается ли качество модели.

Коротко познакомимся с понятиями, которые используются в цикле обучения. Если хочется сразу увидеть код, переходите к полной реализации цикла оптимизации.

Функция потерь

Когда необученная сеть получает обучающие данные, она, скорее всего, даст неправильный ответ. Функция потерь измеряет, насколько полученный результат отличается от целевого значения; именно функцию потерь мы и хотим минимизировать при обучении. Чтобы вычислить потери, мы делаем предсказание по входам из примера данных и сравниваем его с истинным значением метки.

Распространённые функции потерь: nn.MSELoss (среднеквадратичная ошибка, Mean Square Error) для задач регрессии и nn.NLLLoss (отрицательное логарифмическое правдоподобие, Negative Log Likelihood) для классификации. nn.CrossEntropyLoss объединяет nn.LogSoftmax и nn.NLLLoss.

Мы передаём выходные логиты модели в nn.CrossEntropyLoss: она нормализует логиты и вычислит ошибку предсказания.

# Создаём функцию потерь
loss_fn = nn.CrossEntropyLoss()

(Примечание переводчика: поэтому в конце модели нет слоя Softmax — он уже «спрятан» внутри CrossEntropyLoss. Если добавить его ещё и в модель, softmax применится дважды, и обучение пойдёт хуже. Подробно — урок «Правдоподобие, softmax и cross-entropy».)

Оптимизатор

Оптимизация — это процесс подстройки параметров модели, чтобы уменьшить ошибку модели на каждом шаге обучения. Алгоритмы оптимизации определяют, как именно идёт этот процесс (в этом примере мы используем стохастический градиентный спуск, стохастический градиентный спускстохастический градиентный спуск (SGD)EN stochastic gradient descent (SGD)Градиент считают не по всем данным, а по случайному примеру или мини-батчу. В PyTorch torch.optim.SGD — это просто «шаг против градиента», какой бы батч вы ни подали.). Вся логика оптимизации заключена в объекте optimizer. Здесь мы берём оптимизатор SGD; кроме него, в PyTorch есть много других оптимизаторов, например ADAM и RMSProp, которые лучше работают для разных видов моделей и данных.

Мы создаём оптимизатор, регистрируя в нём параметры модели, которые нужно обучать, и передавая гиперпараметр шага обучения.

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

Внутри цикла обучения оптимизация проходит в три шага:

• Вызываем optimizer.zero_grad(), чтобы сбросить градиенты параметров модели. По умолчанию градиенты складываются; чтобы не посчитать их дважды, мы явно обнуляем их на каждой итерации.
• Распространяем ошибку предсказания назад вызовом loss.backward(). PyTorch записывает градиенты функции потерь по каждому параметру.
• Получив градиенты, вызываем optimizer.step(), чтобы подправить параметры на основе градиентов, собранных при обратном проходе.

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

Полная реализация

Определим функцию train_loop, которая выполняет в цикле наш код оптимизации, и test_loop, которая оценивает качество модели на тестовых данных.

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    # Переводим модель в режим обучения — это важно для слоёв батч-нормализации и dropout.
    # Здесь это не обязательно, но так принято делать.
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        # Вычисляем предсказание и потери
        pred = model(X)
        loss = loss_fn(pred, y)

        # Обратное распространение ошибки
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * batch_size + len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn):
    # Переводим модель в режим оценки — это важно для слоёв батч-нормализации и dropout.
    # Здесь это не обязательно, но так принято делать.
    model.eval()
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    # Оценка модели внутри torch.no_grad() гарантирует, что в режиме тестирования градиенты не вычисляются.
    # Заодно это убирает лишние вычисления градиентов и экономит память для тензоров с requires_grad=True.
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

Создаём функцию потерь и оптимизатор и передаём их в train_loop и test_loop. Можете увеличить число эпох и посмотреть, как качество модели растёт.

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
print("Done!")

Вывод (первая эпоха полностью, дальше — только итоги эпох):

Epoch 1
-------------------------------
loss: 2.301126  [   64/60000]
loss: 2.286441  [ 6464/60000]
loss: 2.267145  [12864/60000]
loss: 2.268225  [19264/60000]
loss: 2.244242  [25664/60000]
loss: 2.211982  [32064/60000]
loss: 2.223909  [38464/60000]
loss: 2.181747  [44864/60000]
loss: 2.174136  [51264/60000]
loss: 2.149592  [57664/60000]
Test Error:
 Accuracy: 39.6%, Avg loss: 2.143499

Epoch 2  ...  Accuracy: 51.9%, Avg loss: 1.857896
Epoch 3  ...  Accuracy: 59.8%, Avg loss: 1.495131
Epoch 4  ...  Accuracy: 63.6%, Avg loss: 1.242321
Epoch 5  ...  Accuracy: 65.1%, Avg loss: 1.084806
Epoch 6  ...  Accuracy: 66.0%, Avg loss: 0.981833
Epoch 7  ...  Accuracy: 67.3%, Avg loss: 0.911078
Epoch 8  ...  Accuracy: 68.5%, Avg loss: 0.859727
Epoch 9  ...  Accuracy: 69.7%, Avg loss: 0.820605
Epoch 10 ...  Accuracy: 71.0%, Avg loss: 0.789371
Done!

(Примечание переводчика: строки эпох 2–10 мы сократили; в оригинале каждая эпоха напечатана полностью, как первая. На опубликованной странице весь скрипт выполнился примерно за 2 минуты. Обратите внимание на начальные потери ≈2.30: это $\ln 10$ — столько даёт модель, которая угадывает все 10 классов с равной вероятностью. Такая проверка «первого шага» описана в уроке «Как отлаживать обучение».)

Что почитать дальше

• Функции потерь
• torch.optim
• Warmstart Training a Model

---

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

В этом разделе посмотрим, как сохранять состояние модели, загружать его и запускать модель для предсказаний.

import torch
import torchvision.models as models

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

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

model = models.vgg16(weights='IMAGENET1K_V1')
torch.save(model.state_dict(), 'model_weights.pth')

(Примечание переводчика: эта строка скачивает предобученную сеть VGG16 — около 530 МБ. Чтобы просто попробовать сохранение, достаточно своей маленькой NeuralNetwork из предыдущих разделов: код тот же.)

Чтобы загрузить веса модели, нужно сначала создать экземпляр той же модели, а затем загрузить параметры методом load_state_dict().

В коде ниже мы указываем weights_only=True, чтобы при распаковке (unpickling) выполнялись только те функции, которые нужны для загрузки весов. Использовать weights_only=True при загрузке весов считается хорошей практикой.

model = models.vgg16() # не указываем ``weights``, то есть создаём необученную модель
model.load_state_dict(torch.load('model_weights.pth', weights_only=True))
model.eval()

Заметка

Обязательно вызывайте метод model.eval() перед инференсом, чтобы перевести слои dropout и батч-нормализации в режим оценки. Если этого не сделать, результаты инференса будут непостоянными.

Сохранение и загрузка модели вместе со структурой

При загрузке весов нам сначала пришлось создать экземпляр класса модели, потому что именно класс задаёт структуру сети. Можно сохранить структуру этого класса вместе с моделью — для этого в функцию сохранения передают model (а не model.state_dict()):

torch.save(model, 'model.pth')

Затем модель можно загрузить так, как показано ниже.

Как сказано в документе Saving and loading torch.nn.Modules, лучшей практикой считается сохранять state_dict. Но ниже мы используем weights_only=False, потому что здесь загружается вся модель — это устаревший сценарий использования torch.save.

model = torch.load('model.pth', weights_only=False)

Заметка

Этот способ при сериализации модели использует модуль Python pickle, поэтому при загрузке модели должно быть доступно настоящее определение её класса.

(Примечание переводчика: weights_only=False разрешает pickle выполнить произвольный код из файла. Никогда не загружайте так файлы из непроверенных источников. Готовые веса открытых моделей сегодня обычно распространяют в формате safetensors, где исполняемого кода нет вовсе.)

Связанные туториалы

• Saving and Loading a General Checkpoint in PyTorch
• Tips for loading an nn.Module from a checkpoint

---

Что сделать руками

Это упражнения от переводчика, в оригинале их нет.

1. Запустите полный код из раздела 6 на CPU и запишите точность после каждой эпохи. Затем замените SGD на torch.optim.Adam(model.parameters(), lr=1e-3) и сравните.
2. Посчитайте число параметров модели строкой sum(p.numel() for p in model.parameters()) и сверьте с подсчётом из части 3.
3. Сохраните state_dict после обучения, перезапустите Python, загрузите веса в новую NeuralNetwork и убедитесь, что test_loop даёт ту же точность. Потом уберите model.eval() — в этой модели нет dropout и батч-нормализации, поэтому результат не изменится. Добавьте nn.Dropout(0.2) после первого ReLU и повторите опыт.
4. Перенесите обучение на видеокарту: добавьте device и .to(device) для модели и для X, y, как в быстром старте.

Словарик

EN	RU
hyperparameter	гиперпараметр
learning rate	шаг обучения (скорость обучения)
optimization loop	цикл оптимизации
train loop, test loop	цикл обучения, цикл тестирования
loss function	функция потерь
Mean Square Error	среднеквадратичная ошибка
Negative Log Likelihood	отрицательное логарифмическое правдоподобие
Stochastic Gradient Descent	стохастический градиентный спуск
converge	сходиться
inference	инференс (применение обученной модели)
serialize, unpickling	сериализовать, распаковка pickle