Escribir AlexNet desde Cero en PyTorch

Introducción

Este post es una continuación de la serie sobre la construcción de las redes neuronales convolucionales más populares desde cero en PyTorch. Puede ver el post anterior aquí, donde construimos LeNet5. En este post, construiremos AlexNet, una de las mayores innovaciones clave en la computación visual.

Empezaremos investigando y comprendiendo la arquitectura de AlexNet. A continuación, cargaremos nuestro conjunto de datos, CIFAR-10, y aplicaremos unos procesamientos previos a los datos. A continuación, construiremos nuestro AlexNet desde cero utilizando PyTorch y lo entrenaremos en nuestros datos procesados. Finalmente, la modelo entrenado se evaluará en datos no vistos (pruebas) para fines de evaluación.

Prerrequisitos

Un conocimiento de las redes neuronales será útil para entender este artículo. Esto incluiría estar familiarizado con las diferentes capas de las redes neuronales (capa de entrada, capas ocultas, capa de salida), funciones de activación, algoritmos de optimización (variantes del descenso de gradiente), funciones de pérdida, etc. Además, una familiaridad con la sintaxis de Python y la biblioteca PyTorch es imprescindible para entender los fragmentos de código presentados en este artículo.

Una comprensión de las CNN es fundamental. Esto incluye el conocimiento de las capas convolucionales, las capas de pooling y su papel en la extracción de características de los datos de entrada. Entender conceptos como la tasa de propagación, el relleno y el impacto de la talla del kernel/filtro es también beneficioso.

AlexNet

AlexNet es una red neuronal convolucional profunda, que fue desarrollada inicialmente por Alex Krizhevsky y sus colegas en 2012. Fue diseñada para clasificar imágenes para la competición ImageNet LSVRC-2010, donde logró resultados de Referencia. Puede leer detalladamente sobre el modelo en el documento de investigación original aquí.

Vamos a revisar los puntos clave del artículo de AlexNet. En primer lugar, AlexNet operó con imágenes de 3 canales que tenían un tamaño de (224x224x3). Usó pooling máximo junto con activaciones de ReLU al hacer la subsampling. Los kernel usados para las convulsiones eran de 11×11, 5×5 o 3×3, mientras que los kernel usados para el pooling máximo tenían un tamaño de 3×3. Clasificó imágenes en 1000 clases. También utilizó múltiples GPUs.

Conjunto de datos

Vamos a comenzar cargando y preprocesando los datos. Para nuestros propósitos, utilizaremos el conjunto de datos CIFAR-10. El conjunto de datos consta de 60000 imágenes de 32×32 colores en 10 clases, con 6000 imágenes por clase. Hay 50000 imágenes de entrenamiento y 10000 imágenes de prueba.

Aquí están las clases en el conjunto de datos, así como 10 imágenes de muestra aleatorias de cada una:

Fuente: source

Las clases son completamente mutuamente excluyentes. No hay superposición entre automóviles y camiones. “Automóvil” incluye sedanes, SUVs y cosas de ese tipo. “Camión” incluye solo camiones grandes. Ninguno incluye camionetas.

Importando las Bibliotecas

Vamos a empezar importando las bibliotecas necesarias junto con la definición de una variable device, de manera que el cuaderno sepa utilizar una GPU para entrenar el modelo si está disponible.

import numpy as np
import torch
import torch.nn as nn
from torchvision import datasets
from torchvision import transforms
from torch.utils.data.sampler import SubsetRandomSampler

# Configuración del dispositivo
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Cargando el Conjunto de Datos

Usando torchvision (una biblioteca auxiliar para tareas de visión computacional), cargaremos nuestro conjunto de datos. Este método tiene algunas funciones auxiliares que hacen que la preprocesamiento sea bastante fácil y directo. Vamos a definir las funciones get_train_valid_loader y get_test_loader, y luego llamarlas para cargar y procesar nuestros datos de CIFAR-10:

def get_train_valid_loader(data_dir,
                           batch_size,
                           augment,
                           random_seed,
                           valid_size=0.1,
                           shuffle=True):
    normalize = transforms.Normalize(
        mean=[0.4914, 0.4822, 0.4465],
        std=[0.2023, 0.1994, 0.2010],
    )

    # define transforms
    valid_transform = transforms.Compose([
            transforms.Resize((227,227)),
            transforms.ToTensor(),
            normalize,
    ])
    if augment:
        train_transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ])
    else:
        train_transform = transforms.Compose([
            transforms.Resize((227,227)),
            transforms.ToTensor(),
            normalize,
        ])

    # cargar el conjunto de datos
    train_dataset = datasets.CIFAR10(
        root=data_dir, train=True,
        download=True, transform=train_transform,
    )

    valid_dataset = datasets.CIFAR10(
        root=data_dir, train=True,
        download=True, transform=valid_transform,
    )

    num_train = len(train_dataset)
    indices = list(range(num_train))
    split = int(np.floor(valid_size * num_train))

    if shuffle:
        np.random.seed(random_seed)
        np.random.shuffle(indices)

    train_idx, valid_idx = indices[split:], indices[:split]
    train_sampler = SubsetRandomSampler(train_idx)
    valid_sampler = SubsetRandomSampler(valid_idx)

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=batch_size, sampler=train_sampler)

    valid_loader = torch.utils.data.DataLoader(
        valid_dataset, batch_size=batch_size, sampler=valid_sampler)

    return (train_loader, valid_loader)


def get_test_loader(data_dir,
                    batch_size,
                    shuffle=True):
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    )

    # define transform
    transform = transforms.Compose([
        transforms.Resize((227,227)),
        transforms.ToTensor(),
        normalize,
    ])

    dataset = datasets.CIFAR10(
        root=data_dir, train=False,
        download=True, transform=transform,
    )

    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle
    )

    return data_loader

# conjunto de datos CIFAR10 
train_loader, valid_loader = get_train_valid_loader(data_dir = './data',                                      batch_size = 64,
                       augment = False,                                          random_seed = 1)

test_loader = get_test_loader(data_dir = './data',
                              batch_size = 64)

Vamos a desglosar el código:

• Definimos dos funciones get_train_valid_loader y get_test_loader para cargar respectivamente los conjuntos de entrenamiento/validación y prueba
• Empezamos definiendo la variable normalize con la media y desviación estándar de cada uno de los canales (rojo, verde y azul) del conjunto de datos. Estos pueden calcularse manualmente, pero también están disponibles en línea ya que CIFAR-10 es bastante popular
• Para nuestro conjunto de datos de entrenamiento, agregamos la opción de ampliar el conjunto de datos como well, lo que hace que el entrenamiento sea más robusto y aumenta el número de imágenes también. Nota: la ampliación solo se aplica al subconjunto de entrenamiento y no a los subconjuntos de validación y prueba ya que solo se utilizan para fines de evaluación
• Dividimos el conjunto de datos de entrenamiento en conjuntos de entrenamiento y validación (proporción 90:10), y seleccionamos aleatoriamente un subconjunto de él del conjunto de entrenamiento completo
• Especificamos el tamaño de lote y mezclamos el conjunto de datos al cargar, de modo que cada lote tenga alguna variación en los tipos de etiquetas que tiene. Esto aumentará la eficacia de nuestro modelo resultante.
• Finalmente, hacemos uso de los cargadores de datos. Esto puede no afectar el rendimiento en el caso de un pequeño conjunto de datos como CIFAR-10, pero realmente puede impedir el rendimiento en caso de grandes conjuntos de datos y generalmente se considera una buena práctica. Los cargadores de datos nos permiten iterar por los datos en lotes, y los datos se carguen mientras se itera y no todos a la vez en el inicio de su RAM

AlexNet desde Cero

Vamos a empezar con el código primero:

class AlexNet(nn.Module):
    def __init__(self, num_classes=10):
        super(AlexNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=0),
            nn.BatchNorm2d(96),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 3, stride = 2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 3, stride = 2))
        self.layer3 = nn.Sequential(
            nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(384),
            nn.ReLU())
        self.layer4 = nn.Sequential(
            nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(384),
            nn.ReLU())
        self.layer5 = nn.Sequential(
            nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 3, stride = 2))
        self.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(9216, 4096),
            nn.ReLU())
        self.fc1 = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU())
        self.fc2= nn.Sequential(
            nn.Linear(4096, num_classes))

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.layer5(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        out = self.fc1(out)
        out = self.fc2(out)
        return out

Definiendo el Modelo AlexNet

Vamos a adentrarnos en cómo funciona el código anterior:

• El primer paso para definir cualquier red neuronal (ya sea una CNN o no) en PyTorch es definir una clase que hereda de nn.Module ya que contiene muchos de los métodos que necesitaremos utilizar
• Después de eso, hay dos pasos principales. El primero es inicializar las capas que vamos a usar en nuestra CNN dentro de __init__, y el otro es definir la secuencia en la que esas capas procesarán la imagen. Esto se define dentro de la función forward.
• Para la arquitectura en sí, primero definimos las capas convolucionales usando la función nn.Conv2D con el tamaño de kernel apropiado y los canales de entrada/salida. También aplicamos la max pooling usando la función nn.MaxPool2D. Lo bueno de PyTorch es que podemos combinar la capa convolucional, función de activación y max pooling en una sola capa (se aplicarán por separado, pero ayuda con la organización) usando la función nn.Sequential
• Luego definimos las capas completamente conectadas usando funciones lineales (nn.Linear), dropout (nn.Dropout) y la función de activación ReLU (nn.ReLU), combinándolas con la función nn.Sequential
• Finalmente, nuestra última capa emite 10 neuronas, que son nuestras predicciones finales para las 10 clases de objetos

Estableciendo Hiperparámetros

Antes de entrenar, necesitamos establecer algunos hiperparámetros, como la función de pérdida y el optimizador a usar, junto con el tamaño de lote, la tasa de aprendizaje y el número de épocas.

num_classes = 10
num_epochs = 20
batch_size = 64
learning_rate = 0.005

model = AlexNet(num_classes).to(device)

# Pérdida y optimizador
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, weight_decay = 0.005, momentum = 0.9)  

# Entrenar el modelo
total_step = len(train_loader)

Empezamos definiendo hyperparámetros sencillos (épocas, tamaño de lote y tasa de aprendizaje) e inicializando nuestro modelo usando el número de clases como argumento, que en este caso es 10, junto con transferir el modelo al dispositivo apropiado (CPU o GPU). Luego definimos nuestra función de costo como pérdida de entropía cruzada y el optimizador como Adam. Existen muchas opciones para estos, pero éstos tendrían tendencia a dar buenos resultados con el modelo y los datos dados. Finalmente, definimos total_step para mantener un mejor seguimiento de los pasos durante el entrenamiento

Entrenamiento

En este punto estamos listos para entrenar nuestro modelo:

total_step = len(train_loader)

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):  
        # Movimiento de tensores al dispositivo configurado
        images = images.to(device)
        labels = labels.to(device)

        # Paso hacia delante
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Paso atrás y optimización
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

    # Validation
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in valid_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            del images, labels, outputs

        print('Accuracy of the network on the {} validation images: {} %'.format(5000, 100 * correct / total)) 

Veamos qué hace el código:

• Empezamos iterando por el número de épocas, y luego los lotes en nuestros datos de entrenamiento
• Convertimos las imágenes y las etiquetas según el dispositivo que estamos utilizando, es decir, GPU o CPU
• En el paso hacia delante, hacemos predicciones usando nuestro modelo y calculamos la pérdida basada en esas predicciones y nuestras etiquetas reales
• A continuación, hacemos el paso atrás donde actualmente actualizamos nuestros pesos para mejorar nuestro modelo
• Luego establecemos los gradientes en cero antes de cada actualización usando la función optimizer.zero_grad()
• A continuación, calculamos los nuevos gradientes usando la función loss.backward()
• Y finalmente, actualizamos los pesos con la función optimizer.step()
• También, al final de cada época, usamos nuestro conjunto de validación para calcular la precisión del modelo. En este caso, no necesitamos gradientes así que usamos with torch.no_grad() para una evaluación más rápida

Podemos ver la salida como sigue:

Pérdida de Entrenamiento y Precisión deValidación

Como podemos ver, la pérdida disminuye con cada época, lo que muestra que nuestro modelo realmente está aprendiendo. Tenga en cuenta que esta pérdida es en el conjunto de entrenamiento y si la pérdida es demasiado pequeña, puede indicar sobreajuste. Esta es la razón por la que también usamos el conjunto de validación. La precisión parece incrementar en el conjunto de validación, lo que indica que hay una probabilidad muy baja de cualquier sobreajuste. Vamos a probar nuestro modelo para ver cómo se comporta.

Prueba

Ahora veremos cómo se comporta nuestro modelo con datos no vistos antes:

with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        del images, labels, outputs

    print('Accuracy of the network on the {} test images: {} %'.format(10000, 100 * correct / total))   

Tenga en cuenta que el código es exactamente el mismo que para nuestros propósitos de validación.

Usando el modelo y entrenando solo para 6 épocas, parece que obtenemos alrededor de un 78,8% de precisión en el conjunto de validación.

Precisión de Prueba

Conclusión

Vamos a concluir lo que hicimos en este artículo:

• Empezamos comprendiendo la arquitectura y los diferentes tipos de capas del modelo AlexNet
• A continuación, cargamos y preprocesamos el conjunto de datos CIFAR-10 utilizando torchvision
• Después, utilizamos PyTorch para construir nuestro modelo AlexNet de cero
• Finalmente, entrenamos y probamos nuestro modelo en el conjunto de datos CIFAR-10, y el modelo pareció funcionar bien en el conjunto de pruebas con un entrenamiento mínimo (6 épocas)

Trabajo futuro

Este artículo ofrece una buena introducción y experiencia práctica, pero aún podrás adquirir más conocimientos explorando más allá y descubriendo qué otras cosas puedes lograr.

• Puedes intentar utilizar diferentes conjuntos de datos. Uno de estos es CIFAR-100, que es una extensión del conjunto de datos CIFAR-10 con 100 clases
• Puedes experimentar con diferentes hiperparámetros y ver la mejor combinación de ellos para el modelo
• Finalmente, puedes intentar agregar o eliminar capas del conjunto de datos para ver su impacto en la capacidad del modelo.