Módulo de visión artificial de PyTorch torchvision

torchvision es una biblioteca de extensión en el ecosistema de PyTorch especializada en tareas de visión artificial; proporciona las siguientes funciones principales:

  1. Modelos preentrenados: incluye implementaciones de arquitecturas CNN clásicas (como ResNet, VGG, AlexNet, etc.)
  2. Herramientas de conjuntos de datos: incluye conjuntos de datos visuales comunes (como CIFAR10, MNIST, ImageNet, etc.)
  3. Transformaciones de imágenes: proporciona varios métodos de preprocesamiento de imágenes y aumento de datos
  4. Utilidades: incluye funciones auxiliares como procesamiento de videos y operaciones de imágenes
# 安装 torchvision（通常与 PyTorch 一起安装）
pip install torch torchvision

Análisis de componentes clave

1. torchvision.models

Proporciona modelos de visión artificial preentrenados que se pueden usar directamente para el aprendizaje por transferencia:

Ejemplo

import torchvision.models as models

# Cargar el modelo preentrenado
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
vgg16 = models.vgg16(pretrained=True)

Lista de modelos comunes:

Nombre del modelo Escenarios de aplicación Número de parámetros Precisión Top-1
ResNet Clasificación general de imágenes 11M-60M 69%-80%
VGG Extracción de características 138M 71.3%
MobileNet Aplicaciones móviles 3.4M 70.6%
EfficientNet Modelos eficientes 5M-66M 77%-84%

2. torchvision.datasets

Incluye conjuntos de datos comunes de visión artificial, lo que simplifica el proceso de carga de datos:

Ejemplo

from torchvision import datasets

# Cargar el conjunto de datos CIFAR10
train_data = datasets.CIFAR10(
    root='data',
    train=True,
    download=True,
    transform=transforms.ToTensor()
)

# Cargar el conjunto de datos MNIST
test_data = datasets.MNIST(
    root='data',
    train=False,
    download=True
)

Tipos de conjuntos de datos compatibles:

Ejemplo

graph TD
A[torchvision.datasets] --> B[conjuntos de datos de clasificación]
A --> C[conjuntos de datos de detección]
A --> D[conjuntos de datos de segmentación]
    B --> B1[CIFAR10/100]
    B --> B2[MNIST/FashionMNIST]
    B --> B3[ImageNet]
    C --> C1[COCO]
    C --> C2[VOC]
    D --> D1[Cityscapes]

3. torchvision.transforms

Herramientas principales para el preprocesamiento de imágenes y el aumento de datos:

Ejemplo

from torchvision import transforms

# Definir el pipeline de transformación de imágenes
transform = transforms.Compose([
    transforms.Resize(256),          # Redimensionar
    transforms.CenterCrop(224),       # Recorte central
    transforms.ToTensor(),           # Convertir a tensor
    transforms.Normalize(             # Normalizar
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

Clasificación de métodos de transformación comunes:

Categoría Ejemplos de métodos Función
Transformaciones geométricas RandomRotation, RandomResizedCrop Aumentar la invariancia de posición
Transformaciones de color ColorJitter, Grayscale Mejorar la robustez del color
Desenfoque/Ruido GaussianBlur, RandomErasing Prevenir el sobreajuste
Transformaciones combinadas RandomApply, RandomChoice Estrategias de combinación flexibles

Ejemplo práctico: flujo de clasificación de imágenes

1. Preparación de datos

Ejemplo

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

# Definir las transformaciones de datos
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Cargar el conjunto de datos
train_set = datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=train_transform
)

# Crear el cargador de datos
train_loader = DataLoader(
    train_set,
    batch_size=32,
    shuffle=True
)

2. Entrenamiento del modelo

Ejemplo

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

# Usar el modelo preentrenado
model = models.resnet18(pretrained=True)

# Modificar la última capa (adaptada a las 10 clases de CIFAR10)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)

# Definir la función de pérdida y el optimizador
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Bucle de entrenamiento
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()

Funciones avanzadas

1. Conjuntos de datos personalizados

Ejemplo

from torchvision.datasets import VisionDataset

class CustomDataset(VisionDataset):
    def __init__(self, root, transform=None):
        super().__init__(root, transform=transform)
        # Implementar __getitem__ y __len__
       
    def __getitem__(self, index):
        # Devolver la tupla (image, target)
        pass
       
    def __len__(self):
        # Devolver el tamaño del conjunto de datos
        pass

2. Exportación y despliegue del modelo

Ejemplo

# Exportar al formato ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"]
)

Recomendaciones de mejores prácticas

1. Estrategias de aumento de datos：

  • Usar transformaciones aleatorias para aumentar los datos durante el entrenamiento
  • Para validación/pruebas, usar solo transformaciones deterministas

2. Técnicas de aprendizaje por transferencia：

Ejemplo

# Congelar todos los parámetros excepto la última capa
for param in model.parameters():
    param.requires_grad = False
model.fc.requires_grad = True

3. Optimización del rendimiento：

  • Usarnum_workersel parámetro para acelerar la carga de datos
  • Para conjuntos de datos grandes, considere usarDatasetun subconjunto de

4. Errores comunes：

  • Olvidar llamarzero_grad()
  • Confundirtrain()yeval()modo
  • La forma del tensor de imagen no cumple con los requisitos del modelo (debe ser C×H×W)
Otras extensiones