Обучение нейросети PyTorch: полный гайд по созданию и тренировке моделей

Почему PyTorch стал стандартом для обучения нейросетей

PyTorch — это фреймворк машинного обучения с открытым исходным кодом, созданный на базе библиотеки Torch. Он завоевал популярность благодаря гибкости, динамическому вычислительному графу и тесной интеграции с Python. В отличие от статичных графов, которые использовались в ранних версиях TensorFlow, PyTorch позволяет изменять архитектуру сети на лету, что особенно удобно при исследовательской работе и отладке.

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

Для новичков важно понимать, что PyTorch — это не просто инструмент, а целая экосистема. В неё входят torchvision для работы с изображениями, torchtext для текстов, torchaudio для звука, а также дополнительные библиотеки, такие как Ignite, которые упрощают организацию процесса обучения.

Установка PyTorch и подготовка окружения

Перед началом работы необходимо установить PyTorch. Официальный сайт предлагает конфигуратор, который генерирует команду установки в зависимости от операционной системы, менеджера пакетов и наличия CUDA. Для большинства пользователей достаточно выполнить:

pip install torch torchvision

Если у вас есть видеокарта NVIDIA, рекомендуется установить версию с поддержкой CUDA, чтобы ускорить обучение. Например:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

Важно проверить, что PyTorch видит GPU. Для этого выполните в Python:

import torch
print(torch.cuda.is_available())

Если вы работаете в среде Google Colab или на суперкомпьютере, окружение с PyTorch часто уже предустановлено. В таких случаях достаточно выбрать соответствующий runtime или создать персональное окружение с нужной версией фреймворка.

Также стоит установить Jupyter Notebook или использовать IDE, например VS Code, для удобного написания и запуска кода. Для визуализации процесса обучения пригодятся библиотеки matplotlib и tensorboard.

Основы работы с тензорами в PyTorch

Тензоры — это фундаментальная структура данных в PyTorch, аналогичная массивам NumPy, но с поддержкой автоматического дифференцирования и вычислений на GPU. Создать тензор можно разными способами: из списка, из NumPy-массива, с помощью функций torch.zeros, torch.ones, torch.rand и других.

import torch
import numpy as np

# Из списка
t = torch.tensor([1, 2, 3])

# Из NumPy
arr = np.array([4, 5, 6])
t2 = torch.from_numpy(arr)

# Случайный тензор
rand_t = torch.rand(3, 3)

Тензоры поддерживают множество операций: индексирование, срезы, изменение формы (view, reshape), математические операции (сложение, умножение, матричное умножение), а также тригонометрические и статистические функции. Например, чтобы перемножить матрицы, используйте torch.matmul или оператор @.

Важная особенность — возможность перемещать тензоры между CPU и GPU с помощью метода .to(device). Это позволяет легко масштабировать вычисления:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
t = t.to(device)

Понимание тензоров — обязательный шаг перед изучением нейросетей, так как все данные (входные признаки, веса, градиенты) представлены именно в этом формате.

Создание модели нейросети с помощью nn.Module

В PyTorch модели создаются путём наследования от класса torch.nn.Module. Этот класс предоставляет каркас для определения архитектуры, управления параметрами и автоматического дифференцирования. Минимальная модель выглядит так:

import torch.nn as nn
import torch.nn.functional as F

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

Метод forward определяет, как данные проходят через слои. PyTorch автоматически строит вычислительный граф во время прямого прохода, что позволяет затем вычислять градиенты при обратном распространении.

Для более сложных архитектур можно использовать контейнеры nn.Sequential, nn.ModuleList и nn.ModuleDict. Например, nn.Sequential позволяет объединить слои в последовательность:

model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

Также в PyTorch есть готовые архитектуры в torchvision.models: ResNet, VGG, SqueezeNet и другие. Их можно использовать как есть или дообучать под свою задачу (transfer learning).

Подготовка данных: Dataset и DataLoader

Для эффективного обучения нейросети данные должны быть организованы в батчи. В PyTorch за это отвечают классы torch.utils.data.Dataset и DataLoader. Dataset хранит данные и предоставляет доступ к отдельным элементам, а DataLoader автоматически группирует их в батчи, перемешивает и загружает параллельно.

Создать собственный датасет можно, унаследовавшись от Dataset и реализовав методы __len__ и __getitem__:

from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

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

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

Затем создаётся загрузчик:

from torch.utils.data import DataLoader

dataset = MyDataset(data, labels)
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

Для работы с изображениями удобно использовать torchvision.datasets.ImageFolder, который автоматически читает папки с классами. Также доступны встроенные датасеты, такие как MNIST и CIFAR-10.

Важно правильно настроить трансформации данных (нормализация, аугментация) — это повышает качество модели и снижает переобучение. Например, для CIFAR-10 часто применяют случайные горизонтальные отражения и нормализацию по каналам.

Цикл обучения: функция потерь, оптимизатор и обратное распространение

Обучение нейросети в PyTorch — это итеративный процесс, состоящий из прямого прохода, вычисления функции потерь, обратного распространения ошибки и обновления весов. Базовый цикл выглядит так:

import torch.optim as optim

model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

Функция потерь (criterion) измеряет разницу между предсказаниями и истинными метками. Для классификации часто используют CrossEntropyLoss, для регрессии — MSELoss. Оптимизатор обновляет веса на основе градиентов; популярные варианты — SGD, Adam, RMSprop.

Важно вызывать optimizer.zero_grad() перед каждым шагом, чтобы обнулить градиенты из предыдущей итерации. Иначе они будут накапливаться.

После обучения модель нужно перевести в режим оценки (model.eval()) и отключить вычисление градиентов (with torch.no_grad():), чтобы ускорить инференс и избежать лишних вычислений.

Использование Ignite для упрощения циклов обучения

Хотя написание циклов вручную даёт полный контроль, для типовых задач удобно использовать библиотеку Ignite, которая факторизует процесс обучения и валидации. Ignite предоставляет класс Engine, который управляет итерациями по эпохам и батчам, а также систему событий для добавления обработчиков.

Установка:

pip install pytorch-ignite

Пример создания тренера:

from ignite.engine import Engine, create_supervised_trainer

trainer = create_supervised_trainer(model, optimizer, criterion, device)

Метод trainer.run(train_loader, max_epochs=10) запускает обучение. Для валидации создаётся отдельный движок create_supervised_evaluator, который вычисляет метрики, такие как точность, полнота и loss.

С помощью событий можно добавить логирование, сохранение моделей, раннюю остановку и изменение скорости обучения. Например, чтобы выводить loss каждые 50 итераций:

from ignite.engine import Events

@trainer.on(Events.ITERATION_COMPLETED)
def log_loss(engine):
    if engine.state.iteration % 50 == 0:
        print(f"Epoch {engine.state.epoch}, Iter {engine.state.iteration}, Loss: {engine.state.output}")

Ignite также умеет сохранять лучшие модели и оптимизатор, что удобно для длительного обучения. Эта библиотека — хороший компромисс между гибкостью чистого PyTorch и удобством высокоуровневых API.

Оценка модели и метрики качества

После обучения необходимо оценить качество модели на тестовой выборке. Основные метрики для классификации — accuracy, precision, recall и F1-score. В PyTorch их можно вычислить вручную или с помощью библиотек, таких как Ignite или sklearn.

Пример вычисления accuracy вручную:

correct = 0
total = 0
with torch.no_grad():
    for data, target in test_loader:
        output = model(data)
        _, predicted = torch.max(output.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()

print(f'Accuracy: {100 * correct / total}%')

В Ignite метрики подключаются к движку оценщика:

from ignite.metrics import Accuracy, Loss

metrics = {
    'accuracy': Accuracy(),
    'loss': Loss(criterion)
}
evaluator = create_supervised_evaluator(model, metrics=metrics, device=device)

Важно помнить, что метрики нужно вычислять на данных, которые модель не видела во время обучения, иначе результаты будут завышены. Также стоит следить за переобучением: если accuracy на тренировочной выборке значительно выше, чем на тестовой, модель переобучилась.

Сохранение и загрузка моделей, чекпоинты

После обучения модель нужно сохранить для дальнейшего использования. В PyTorch есть два основных способа: сохранить только веса (state_dict) или всю модель целиком. Рекомендуется сохранять state_dict, так как это более гибко и совместимо с изменениями архитектуры.

Сохранение:

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

Загрузка:

model = SimpleNet()
model.load_state_dict(torch.load('model.pth'))
model.eval()

Для длительного обучения полезно сохранять чекпоинты, включающие состояние модели, оптимизатора и номер эпохи. Это позволяет возобновить обучение с того же места в случае сбоя:

checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss
}
torch.save(checkpoint, 'checkpoint.pth')

При загрузке чекпоинта нужно восстановить все компоненты. В Ignite есть встроенные обработчики для сохранения лучших моделей, что упрощает процесс.

Практические советы и типичные ошибки

Начинающие часто сталкиваются с рядом типичных проблем при обучении нейросетей. Вот несколько советов, которые помогут их избежать.

1. Не забывайте про model.train() и model.eval(). Эти методы переключают режимы dropout и batch normalization, что критически важно для корректного обучения и инференса.

2. Правильно нормализуйте данные. Нормализация входных признаков ускоряет сходимость и улучшает качество. Для изображений обычно используют среднее и стандартное отклонение по каналам.

3. Используйте аугментацию данных. Для задач компьютерного зрения случайные повороты, отражения и изменения яркости помогают модели обобщаться и снижают переобучение.

4. Следите за переобучением. Если loss на тренировке уменьшается, а на валидации растёт, применяйте регуляризацию (L2, dropout) или уменьшайте сложность модели.

5. Выбирайте правильный размер батча. Слишком маленький батч приводит к шумным градиентам, слишком большой — требует много памяти. Экспериментируйте с значениями 16, 32, 64.

6. Используйте GPU. Если доступна видеокарта, обязательно переносите модель и данные на GPU с помощью .to(device). Это ускоряет обучение в разы.

7. Не игнорируйте скорость обучения. Слишком высокая скорость приводит к расходимости, слишком низкая — к медленной сходимости. Используйте планировщики скорости обучения, например StepLR или ReduceLROnPlateau.

Вопросы и ответы