PyTorch Ejemplo - Proyecto de clasificación de imágenes

La clasificación de imágenes es una de las tareas más básicas en la visión por computadora. Su objetivo es permitir que la computadora reconozca el contenido principal de una imagen y lo clasifique en categorías predefinidas. Por ejemplo, reconocer si una imagen contiene un gato o un perro.

Aplicación del aprendizaje profundo en la clasificación de imágenes

Los modelos de aprendizaje profundo, especialmente las redes neuronales convolucionales (CNN), se han convertido en la solución principal para las tareas de clasificación de imágenes. PyTorch, como framework de aprendizaje profundo, proporciona un conjunto completo de herramientas para construir y entrenar modelos CNN.

Descripción general del flujo del proyecto

Un proyecto completo de clasificación de imágenes generalmente incluye los siguientes pasos:

  1. Preparación y preprocesamiento de datos
  2. Construcción del modelo
  3. Entrenamiento del modelo
  4. Evaluación del modelo
  5. Aplicación del modelo

Preparación del entorno y carga de datos

Instalación de las bibliotecas necesarias

# 安装PyTorch和torchvision
!pip install torch torchvision

Carga de conjuntos de datos comunes

El módulo torchvision de PyTorch proporciona varios conjuntos de datos comunes, como CIFAR-10, MNIST, etc.

Ejemplo

import torch
import torchvision
import torchvision.transforms as transforms

# Definir transformaciones de datos
transform = transforms.Compose([
    transforms.ToTensor(),  # Convertir imagen PIL a Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalización
])

# Cargar conjunto de entrenamiento 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)

# Cargar conjunto de prueba 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)

# Definir nombres de categorías
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

Construcción del modelo de red neuronal convolucional

Estructura básica de CNN

Una CNN típica incluye capas convolucionales, capas de agrupación (pooling) y capas totalmente conectadas.

Ejemplo

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # Capa convolucional 1: entrada de 3 canales (RGB), salida de 6 canales, núcleo de convolución 5x5
        self.conv1 = nn.Conv2d(3, 6, 5)
        # Capa de pooling: ventana 2x2, paso 2
        self.pool = nn.MaxPool2d(2, 2)
        # Capa convolucional 2: entrada de 6 canales, salida de 16 canales, núcleo de convolución 5x5
        self.conv2 = nn.Conv2d(6, 16, 5)
        # Capa totalmente conectada 1: entrada 16*5*5, salida 120
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        # Capa totalmente conectada 2: entrada 120, salida 84
        self.fc2 = nn.Linear(120, 84)
        # Capa totalmente conectada 3: entrada 84, salida 10 (correspondiente a 10 categorías)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Primera capa: convolución + ReLU + pooling
        x = self.pool(F.relu(self.conv1(x)))
        # Segunda capa: convolución + ReLU + pooling
        x = self.pool(F.relu(self.conv2(x)))
        # Aplanar mapas de características
        x = x.view(-1, 16 * 5 * 5)
        # Capa totalmente conectada + ReLU
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        # Capa de salida
        x = self.fc3(x)
        return x

# Instanciar la red
net = Net()

Visualización de la estructura del modelo

Entrenamiento del modelo

Definición de la función de pérdida y el optimizador

Ejemplo

import torch.optim as optim

# Función de pérdida de entropía cruzada
criterion = nn.CrossEntropyLoss()
# Optimizador de descenso de gradiente estocástico (SGD)
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

Bucle de entrenamiento

Ejemplo

for epoch in range(10):  # Entrenar durante 10 épocas
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # Obtener datos de entrada
        inputs, labels = data
       
        # Poner a cero los gradientes
        optimizer.zero_grad()
       
        # Propagación hacia adelante
        outputs = net(inputs)
        # Calcular la pérdida
        loss = criterion(outputs, labels)
        # Propagación hacia atrás
        loss.backward()
        # Actualizar pesos
        optimizer.step()
       
        # Imprimir estadísticas
        running_loss += loss.item()
        if i % 2000 == 1999:  # Imprimir cada 2000 mini-lotes
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Finished Training')

Evaluación del modelo

Cálculo de la exactitud en el conjunto de prueba

Ejemplo

correct = 0
total = 0
with torch.no_grad():  # No calcular gradientes
    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'Accuracy on test images: {100 * correct / total:.2f}%')

Análisis de exactitud por categoría

Ejemplo

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'Accuracy of {classes[i]:5s}: {100 * class_correct[i] / class_total[i]:.2f}%')

Guardado y carga del modelo

Guardar el modelo entrenado

Ejemplo

# Guardar parámetros del modelo
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

Cargar el modelo para realizar predicciones

Ejemplo

# Cargar el modelo
net = Net()
net.load_state_dict(torch.load(PATH))

# Usar el modelo para predecir
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}' for j in range(4)))
Otras extensiones