Пример PyTorch - проект классификации изображений

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

Применение глубокого обучения в классификации изображений

Модели глубокого обучения, особенно свёрточные нейронные сети (CNN), стали основным решением для задач классификации изображений. Как фреймворк глубокого обучения, PyTorch предоставляет полный набор инструментов для построения и обучения моделей CNN.

Обзор процесса проекта

Полный проект классификации изображений обычно включает следующие этапы:

  1. Подготовка и предварительная обработка данных
  2. Построение модели
  3. Обучение модели
  4. Оценка модели
  5. Применение модели

Подготовка окружения и загрузка данных

Установка необходимых библиотек

# 安装PyTorch和torchvision
!pip install torch torchvision

Загрузка распространённых наборов данных

torchvision из PyTorch предоставляет несколько распространённых наборов данных, таких как CIFAR-10, MNIST и т. д.

Пример

import torch
import torchvision
import torchvision.transforms as transforms

# Определить преобразования данных
transform = transforms.Compose([
    transforms.ToTensor(),  # Преобразовать изображение PIL в Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Нормализация
])

# Загрузить обучающий набор CIFAR-10
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                       download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                         shuffle=True, num_workers=2)

# Загрузить тестовый набор CIFAR-10
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                      download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                        shuffle=False, num_workers=2)

# Определить имена классов
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'собака', 'лягушка', 'лошадь', 'корабль', 'грузовик')

Построение модели свёрточной нейронной сети

Базовая структура CNN

Типичная CNN состоит из свёрточных слоёв, слоёв пулинга и полносвязных слоёв.

Пример

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # Свёрточный слой 1: вход 3 канала (RGB), выход 6 каналов, ядро 5x5
        self.conv1 = nn.Conv2d(3, 6, 5)
        # Слой пулинга: окно 2x2, шаг 2
        self.pool = nn.MaxPool2d(2, 2)
        # Свёрточный слой 2: вход 6 каналов, выход 16 каналов, ядро 5x5
        self.conv2 = nn.Conv2d(6, 16, 5)
        # Полносвязный слой 1: вход 16*5*5, выход 120
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        # Полносвязный слой 2: вход 120, выход 84
        self.fc2 = nn.Linear(120, 84)
        # Полносвязный слой 3: вход 84, выход 10 (соответствует 10 классам)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Первый свёрточный слой + ReLU + пулинг
        x = self.pool(F.relu(self.conv1(x)))
        # Второй свёрточный слой + ReLU + пулинг
        x = self.pool(F.relu(self.conv2(x)))
        # Развернуть карты признаков
        x = x.представление(-1, 16 * 5 * 5)
        # Полносвязный слой + ReLU
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        # Выходной слой
        x = self.fc3(x)
        return x

# Создание экземпляра сети
net = Net()

Визуализация структуры модели

Обучение модели

Определение функции потерь и оптимизатора

Пример

import torch.optim as optim

# Функция потерь кросс-энтропии
criterion = nn.CrossEntropyLoss()
# Оптимизатор стохастического градиентного спуска
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

Цикл обучения

Пример

for epoch in range(10):  # Обучение на протяжении 10 эпох
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # Получить входные данные
        inputs, labels = data
       
        # Обнулить градиенты
        optimizer.zero_grad()
       
        # Прямое распространение
        outputs = net(inputs)
        # Вычислить потери
        loss = criterion(outputs, labels)
        # Обратное распространение
        loss.backward()
        # Обновить веса
        optimizer.step()
       
        # Вывести статистику
        running_loss += loss.item()
        if i % 2000 == 1999:  # Печатать каждые 2000 мини-пакетов
            print(f'[{epoch + 1}, {i + 1:5d}] потеря: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Обучение завершено')

Оценка модели

Вычисление точности на тестовом наборе

Пример

correct = 0
total = 0
with torch.no_grad():  # Не вычислять градиенты
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Точность на тестовых изображениях: {100 * correct / total:.2f}%')

Анализ точности по классам

Пример

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1

for i in range(10):
    print(f'Точность для {classes[i]:5s}: {100 * class_correct[i] / class_total[i]:.2f}%')

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

Сохранение обученной модели

Пример

# Сохранить параметры модели
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

Загрузка модели для предсказания

Пример

# Загрузить модель
net = Net()
net.load_state_dict(torch.load(PATH))

# Используйте модель для получения прогнозов
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Предсказано: ', ' '.join(f'{classes[predicted[j]]:5s}' for j in range(4)))
Другие расширения