Библиотека PyTorch: нейронные сети на Python

Раздел: Машинное обучение -> Глубокое обучение

Основы работы с PyTorch

PyTorch (torch) - библиотека с открытым исходным кодом для задач глубокого обучения. Основные компоненты: тензоры (аналоги массивов NumPy с поддержкой GPU), автоматическое дифференцирование (autograd) для вычисления градиентов, модули для построения нейронных сетей (nn.Module) и утилиты для загрузки данных.

Рассмотрим построение простой полносвязной сети для классификации изображений (набор MNIST). Сначала импортируем необходимые модули:


import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader

Torch python (фреймворк pytorch)

Загружаем данные и создаём загрузчик:


transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

библиотека torch python (библиотека pytorch в python)

Определяем нейронную сеть:


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

    def forward(self, x):
        x = x.view(-1, 784)  # из [batch,1,28,28] в [batch,784]
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = SimpleNet()

классификация изображений python с использованием resnet50 (классификация изображений с resnet50)

Задаём функцию потерь и оптимизатор:


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

Цикл обучения (одна эпоха):


for epoch in range(1):
    model.train()
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # Оценка на тестовой выборке
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f'Accuracy: {100 * correct / total:.2f}%')

Этот код составляет основу обучения любой нейронной сети в PyTorch.

Типичные ошибки и их решение:

  • Ошибка размерности: при передаче изображений без преобразования в одномерный вектор. Решение: использовать x.view(-1, 784) или nn.Flatten().
  • Потери не уменьшаются: возможно, слишком большой или малый learning rate. Стоит проверить значения лосса на первых итерациях.
  • Ошибка CUDA out of memory: уменьшить batch_size или использовать меньшую модель.

Как создать нейронную сеть без использования nn.Module?

Можно определить слои вручную и использовать только функционал nn.functional, но такой подход менее нагляден и редко применяется на практике.


import torch.nn.functional as F

class SimpleNetManual(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1_weight = nn.Parameter(torch.randn(784, 128) * 0.1)
        self.fc1_bias = nn.Parameter(torch.zeros(128))
        self.fc2_weight = nn.Parameter(torch.randn(128, 64) * 0.1)
        self.fc2_bias = nn.Parameter(torch.zeros(64))
        self.fc3_weight = nn.Parameter(torch.randn(64, 10) * 0.1)
        self.fc3_bias = nn.Parameter(torch.zeros(10))

    def forward(self, x):
        x = x.view(-1, 784)
        x = F.relu(F.linear(x, self.fc1_weight, self.fc1_bias))
        x = F.relu(F.linear(x, self.fc2_weight, self.fc2_bias))
        x = F.linear(x, self.fc3_weight, self.fc3_bias)
        return x

Проблема:

Неудобно управлять весами, проще использовать nn.Linear.

Как использовать кастомную функцию потерь?

Можно создать собственный класс, наследующий nn.Module, или простую функцию.


def mae_loss(output, target):
    return torch.mean(torch.abs(output - target))

# Использование в цикле
loss = mae_loss(predictions, targets)

Ошибка:

Необходимо обеспечить дифференцируемость функции. Использовать операции из torch, а не из Python.

Как переключиться на GPU?


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

for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)
    # ... остальной код

Типичная ошибка:

Если данные остались на CPU, модель не сможет обработать их - произойдёт исключение. Проверять images.device == model.device.

Как сохранить и загрузить модель?


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

# Загрузка
model = SimpleNet()
model.load_state_dict(torch.load('model.pth', map_location=device))
model.eval()

Распространённая ошибка:

Загружать state_dict в модель с другой архитектурой. Keys должны совпадать.

Расширенные примеры работы с PyTorch

Использование nn.Sequential для создания модели

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

Пример 

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

Результат: создаётся тот же классификатор, что и в примере выше.

Пользовательский модуль с несколькими слоями и пропуском соединения (skip connection)

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

Пример 

class ResidualBlock(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear1 = nn.Linear(in_features, out_features)
        self.bn1 = nn.BatchNorm1d(out_features)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(out_features, out_features)
        self.bn2 = nn.BatchNorm1d(out_features)
        self.shortcut = nn.Identity() if in_features == out_features else nn.Linear(in_features, out_features)

    def forward(self, x):
        identity = self.shortcut(x)
        out = self.relu(self.bn1(self.linear1(x)))
        out = self.bn2(self.linear2(out))
        out += identity
        return self.relu(out)

class ResNetLike(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.block1 = ResidualBlock(256, 256)
        self.block2 = ResidualBlock(256, 128)
        self.fc2 = nn.Linear(128, 10)
        self.flatten = nn.Flatten()

    def forward(self, x):
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.fc2(x)
        return x

model = ResNetLike()
print(model)

Проблема:

Неправильные размерности при пропуске. Перед сложением нужно убедиться, что shortcut выдаёт тензор той же формы, что и основной поток.

Использование планировщика скорости обучения (scheduler)

Планировщики изменяют learning rate во время обучения, что часто улучшает сходимость.

Пример 

scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

for epoch in range(10):
    # обучение...
    scheduler.step()
    print(f"Epoch {epoch}, LR: {optimizer.param_groups[0]['lr']:.6f}")

Epoch 0, LR: 0.010000
Epoch 1, LR: 0.010000
Epoch 2, LR: 0.010000
Epoch 3, LR: 0.001000
...

Трансферное обучение с предобученной моделью (torchvision)

Загружаем предобученный ResNet-18 и заменяем последний полносвязный слой для нового числа классов.

Пример 

import torchvision.models as models

model = models.resnet18(pretrained=True)
# Заморозить все слои, кроме последнего
for param in model.parameters():
    param.requires_grad = False

# Заменяем выходной слой
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)  # для 10 классов

# Тренируем только последний слой
optimizer = optim.SGD(model.fc.parameters(), lr=0.001)

Работа с нестандартными данными: пользовательский Dataset

Пример Dataset, который загружает изображения из папки и применяет аугментацию.

Пример 

from PIL import Image
import os

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.images = [f for f in os.listdir(root_dir) if f.endswith('.png')]

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.root_dir, self.images[idx])
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        label = 0  # пример, можно загрузить из имени файла
        return image, label

Ручное вычисление градиентов без autograd (не рекомендуется)

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

Пример 

x = torch.randn(3, requires_grad=False)
w = torch.randn(3, requires_grad=True)
b = torch.randn(1, requires_grad=True)
y = (x * w).sum() + b

# Принудительно вызвать backward
y.backward()
print(w.grad)  # градиент по w

tensor([ 0.2134, -1.2356,  0.7843])

Использование нескольких оптимизаторов для разных частей модели

Например, для GAN: дискриминатор и генератор обучаются с разными learning rate.

Пример 

optimizerD = optim.Adam(discriminator.parameters(), lr=0.0002)
optimizerG = optim.Adam(generator.parameters(), lr=0.0001)

Параллельное обучение на нескольких GPU (DataParallel)

Пример 

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)
model.to(device)

Проблема:

DataParallel может вызывать медленную работу из-за накладных расходов. Для современных архитектур лучше использовать DistributedDataParallel.

Библиотека PyTorch в Python - comments

En
библиотека torch python (python)