Модуль компьютерного зрения torchvision PyTorch

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

  1. Предварительно обученные модели: включает реализации классических архитектур CNN (таких как ResNet, VGG, AlexNet и т. д.)
  2. Инструменты для наборов данных: встроенные распространённые наборы визуальных данных (такие как CIFAR10, MNIST, ImageNet и т. д.)
  3. Преобразования изображений: предоставляет различные методы предварительной обработки изображений и аугментации данных
  4. Служебные инструменты: включает вспомогательные функции, такие как обработка видео, операции с изображениями и т. д.
# 安装 torchvision（通常与 PyTorch 一起安装）
pip install torch torchvision

Анализ ключевых компонентов

1. torchvision.models

Предоставляет предварительно обученные модели компьютерного зрения, которые можно напрямую использовать для трансферного обучения:

Пример

import torchvision.models as models

# Загрузить предварительно обученную модель
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
vgg16 = models.vgg16(pretrained=True)

Список популярных моделей:

Название модели Применимый сценарий Количество параметров Точность Top-1
ResNet Общая классификация изображений 11M-60M 69%-80%
VGG Извлечение признаков 138M 71.3%
MobileNet Мобильное приложение 3.4M 70.6%
EfficientNet Эффективная модель 5M-66M 77%-84%

2. torchvision.datasets

Встроенные популярные наборы данных компьютерного зрения для упрощения процесса загрузки данных:

Пример

from torchvision import datasets

# Загрузить набор данных CIFAR10
train_data = datasets.CIFAR10(
    root='data',
    train=True,
    download=True,
    transform=transforms.ToTensor()
)

# Загрузить набор данных MNIST
test_data = datasets.MNIST(
    root='data',
    train=False,
    download=True
)

Поддерживаемые типы наборов данных:

Пример

graph TD
A[torchvision.datasets] --> B[Наборы данных для классификации]
A --> C[Наборы данных для детекции]
A --> D[Наборы данных для сегментации]
    B --> B1[CIFAR10/100]
    B --> B2[MNIST/FashionMNIST]
    B --> B3[ImageNet]
    C --> C1[COCO]
    C --> C2[VOC]
    D --> D1[Cityscapes]

3. torchvision.transforms

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

Пример

from torchvision import transforms

# Определить конвейер преобразования изображений
transform = transforms.Compose([
    transforms.Resize(256),          # Изменение размера
    transforms.CenterCrop(224),       # Центральная обрезка
    transforms.ToTensor(),           # Преобразовать в тензор
    transforms.Normalize(             # Нормализация
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

Классификация распространённых методов преобразования:

Категория Примеры методов Назначение
Геометрические преобразования RandomRotation, RandomResizedCrop Увеличение инвариантности к положению
Цветовые преобразования ColorJitter, Grayscale Повышение устойчивости к цветовым изменениям
Размытие/Шум GaussianBlur, RandomErasing Предотвращение переобучения
Комбинированные преобразования RandomApply, RandomChoice Гибкая стратегия комбинирования

Практический пример: конвейер классификации изображений

1. Подготовка данных

Пример

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Определяем преобразования данных
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Загрузить набор данных
train_set = datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=train_transform
)

# Создать загрузчик данных
train_loader = DataLoader(
    train_set,
    batch_size=32,
    shuffle=True
)

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

Пример

import torch.nn as nn
import torch.optim as optim

# Использовать предварительно обученную модель
model = models.resnet18(pretrained=True)

# Изменить последний слой (адаптировать к 10 классам CIFAR10)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)

# Определяем функцию потерь и оптимизатор
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Цикл обучения
for epoch in range(10):
    for images, labels in train_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)
       
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Расширенные функции

1. Пользовательский набор данных

Пример

from torchvision.datasets import VisionDataset

class CustomDataset(VisionDataset):
    def __init__(self, root, transform=None):
        super().__init__(root, transform=transform)
        # Реализуйте __getitem__ и __len__
       
    def __getitem__(self, index):
        # Возвращает кортеж (image, target)
        pass
       
    def __len__(self):
        # Возвращает размер набора данных
        pass

2. Экспорт и развёртывание модели

Пример

# Экспорт в формат ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"]
)

Рекомендации по лучшим практикам

1. Стратегия аугментации данных：

  • Используйте случайные преобразования для аугментации данных во время обучения
  • Используйте только детерминированные преобразования при валидации/тестировании

2. Советы по переносу обучения：

Пример

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

3. Оптимизация производительности：

  • Используйтеnum_workersпараметр для ускорения загрузки данных
  • Для больших наборов данных рассмотрите использованиеDatasetподмножество

4. Распространённые ошибки：

  • Забыли вызватьzero_grad()
  • Перепуталиtrain()иeval()mode
  • Форма тензора изображения не соответствует требованиям модели (должна быть C×H×W)
Другие расширения