Escrevendo AlexNet de Raiz em PyTorch

Introdução

Este post é uma continuação da série sobre a construção das redes neuralizadas convolucionais mais populares a partir do zero no PyTorch. Você pode ver o post anterior aqui, onde construímos LeNet5. Neste post, construiremos AlexNet, uma das Breakthrough Algorithms mais chave na visão computacional.

Começaremos investigando e entendendo a arquitetura de AlexNet. Em seguida, mergulharmos diretamente no código carregando o nosso conjunto de dados, CIFAR-10, antes de aplicarmos algumas pré-processações ao dado. Em seguida, construiremos o nosso AlexNet de raiz usando o PyTorch e o treinaremos nos nossos dados pré-processados. Finalmente, o modelo treinado será testado em dados não vistos (testes) para fins de avaliação.

Pré-requisitos

O conhecimento de redes neurais será útil para entender este artigo. Isso incluiria estar familiarizado com as diferentes camadas de redes neurais (camada de entrada, camadas ocultas, camada de saída), funções de ativação, algoritmos de otimização (variações do descida de gradiente), funções de perda, etc. Adicionalmente, a familiaridade com a sintaxe do Python e com a biblioteca PyTorch é essencial para entender os trechos de código apresentados neste artigo.

Uma compreensão dos CNNs é fundamental. Isso inclui o conhecimento de camadas convolucionais, camadas de pooling, e seu papel na extração de características de dados de entrada. Entender conceitos como stride, padding, e o impacto do tamanho do kernel/filtro é também benéfico.

AlexNet

O AlexNet é uma rede neural profunda convolucional, que foi inicialmente desenvolvida por Alex Krizhevsky e seus colegas em 2012. Foi projetada para classificar imagens para a competição ImageNet LSVRC-2010, onde obteve resultados de ponta. Você pode ler detalhes sobre o modelo no artigo de pesquisa original aqui.

Vamos passar por as principais considerações do artigo do AlexNet. Primeiro, o AlexNet operou com imagens de 3 canais que eram (224x224x3) em tamanho. Usou max pooling juntamente com ativações de ReLU quando subamostrando. Os kernels usados para convoluções eram de 11×11, 5×5 ou 3×3, enquanto os kernels usados para max pooling eram de tamanho 3×3. Classificou imagens em 1000 classes. Também utilizou vários GPUs.

Conjunto de Dados

Vamos começar carregando e depois pré-processando os dados. Para nossos propósitos, vamos usar o conjunto de dados CIFAR-10. O conjunto de dados consiste the 60000 imagens coloridas de 32×32 em 10 classes, com 6000 imagens por classe. Há 50000 imagens de treinamento e 10000 imagens de teste.

Aqui estão as classes no conjunto de dados, bem como 10 imagens de amostra aleatórias de cada classe:

Fonte: source

As classes são mutuamente exclusivas. Não há sobreposição entre automóveis e caminhões. “Automóvel” inclui sedãs, SUVs e coisas do tipo. “Caminhão” inclui apenas caminhões grandes. Nenhum inclui caminhões pick-up.

Importando as Bibliotecas

Vamos começar importando as bibliotecas necessárias juntamente com a definição de uma variável device, para que o bloco de notas saiba usar uma GPU para treinar o modelo, se disponível.

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

# Configuração do Dispositivo
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Carregando o Conjunto de Dados

Usando o torchvision (uma biblioteca de auxílio para tarefas de visão computacional), nós carregaremos nossa base de dados. Este método tem algumas funções auxiliares que tornam a pré-processamento bem fácil e direto. Vamos definir as funções get_train_valid_loader e get_test_loader, e então chamá-las para carregar e processar nossos dados 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 transformaçãoações
    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,
        ])

    # carrega a base de dados
    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 transformação
    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 dados 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 descompactar o código:

• Nós definimos duas funções get_train_valid_loader e get_test_loader para carregar conjuntos de treinamento/validação e teste, respectivamente
• Nós começamos definindo a variável normalize com as médias e desvios padrões de cada canal (vermelho, verde e azul) na base de dados. Esses podem ser calculados manualmente, mas estão também disponíveis online, já que o CIFAR-10 é bastante popular
• Para nosso conjunto de treinamento, nós adicionamos a opção de ampliar o conjunto de dados, assim como torná-lo mais robusto e aumentar o número de imagens. Nota: a ampliação é aplicada apenas ao subconjunto de treinamento e não às subdivisões de validação e teste, já que elas são usadas apenas para fins de avaliação
• Nós dividimos o conjunto de treinamento em treinamento e validação (taxa de divisão de 90:10), e pegamos um subconjunto aleatório do conjunto de treinamento inteiro
• Nós especificamos o tamanho do lote e embaralhamos o conjunto de dados ao carregá-lo, para que cada lote tenha alguma variação nos tipos de rótulos que contém. Isso aumentará a eficácia do modelo resultante.
• Finalmente, nós usamos carregadores de dados. Isto pode não afetar o desempenho no caso de um pequeno conjunto de dados como o CIFAR-10, mas pode realmente impedir o desempenho no caso de grandes conjuntos de dados e é considerada uma boa prática em geral. Carregadores de dados nos permitem iterar pelos dados em lotes, e o dado é carregado enquanto iteramos e não é carregado todo de uma vez no início do RAM

AlexNet de Raiz

Vamos começar com o código primeiro:

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

Definindo o Modelo AlexNet

Vamos mergulhar em como o código acima funciona:

• O primeiro passo para definir qualquer rede neural (seja uma CNN ou não) em PyTorch é definir uma classe que herda de nn.Module pois contém muitos dos métodos que precisaremos de utilizar
• Existem duas etapas principais depois disso. A primeira é inicializar as camadas que vamos usar em nossa CNN dentro de __init__, e a outra é definir a sequência na qual essas camadas processarão a imagem. Isto é definido dentro da função forward.
• Para a arquitetura em si, primeiro definimos as camadas convolucionais usando a função nn.Conv2D com o tamanho de kernel apropriado e os canais de entrada/saída. Também aplicamos o max pooling usando a função nn.MaxPool2D. O que é legal sobre PyTorch é que podemos combinar a camada convolucional, função de ativação e max pooling em uma única camada (elas serão aplicadas separadamente, mas ajuda com a organização) usando a função nn.Sequential
• Então, definimos as camadas totalmente conectadas usando funções lineares (nn.Linear) e dropout (nn.Dropout) juntamente com a função de ativação ReLu (nn.ReLU) e combinando esses com a função nn.Sequential
• Finalmente, nossa última camada exibe 10 neurônios, que são as nossas previsões finais para as 10 classes de objetos

Configurando Hiperparâmetros

Antes de treinar, precisamos definir alguns hiperparâmetros, como a função de perda e o otimizador a serem usados juntamente com o tamanho do lote, taxa de aprendizagem e número de épocas.

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

model = AlexNet(num_classes).to(device)

# Função de perda e otimizador
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, weight_decay = 0.005, momentum = 0.9)  

# Treinando o modelo
total_step = len(train_loader)

Começamos definindo hyperparâmetros simples (épocas, tamanho de lote e taxa de aprendizagem) e inicializando o nosso modelo usando o número de classes como argumento, que neste caso é 10, juntamente com a transferência do modelo para o dispositivo apropriado (CPU ou GPU). Em seguida, definimos nossa função de custo como perda de entropia cruzada e o otimizador como Adam. Existem muitas escolhas para estes, mas essas tendem a dar bons resultados com o modelo e os dados dados. Finalmente, definimos total_step para melhor rastrear os passos durante o treinamento

Treinamento

Nós estamos prontos para treinar o nosso modelo neste ponto:

total_step = len(train_loader)

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):  
        # Move os tensores para o dispositivo configurado
        images = images.to(device)
        labels = labels.to(device)

        # Passo de frente
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Passo de trás e otimização
        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)) 

Vamos ver o que o código faz:

• Começamos iterando pelo número de épocas, e então pelos lotes em nossos dados de treinamento
• Nós convertemos as imagens e as legendas de acordo com o dispositivo que estamos usando, por exemplo, GPU ou CPU
• Na passagem de frente, nós fazemos previsões usando o nosso modelo e calculamos a perda com base nas previsões e nas nossas legendas reais
• A seguir, nós fazemos a passagem de trás onde atualmente nossos pesos para melhorar o nosso modelo
• Então, nós definimos os gradientes para zero antes de cada atualização usando a função optimizer.zero_grad()
• Em seguida, calculamos os novos gradientes usando a função loss.backward()
• E, finalmente, atualizamos os pesos com a função optimizer.step()
• Também, no final de cada época, usamos nosso conjunto de validação para calcular a acurácia do modelo. Neste caso, não precisamos de gradientes, portanto usamos with torch.no_grad() para uma avaliação rápida

Podemos ver a saída da seguinte forma:

Perda de Treinamento e Acurácia de Validação

Como podemos ver, a perda está decrescendo a cada época, o que mostra que o nosso modelo está realmente aprendendo. Note que esta perda é no conjunto de treinamento e, se a perda for muito pequena, isso pode indicar sobreajuste. É por isso que estamos usando o conjunto de validação também. A acurácia parece estar aumentando no conjunto de validação, o que indica que há pouca probabilidade de sobreajuste. Vamos agora testar o nosso modelo para ver como ele performa.

Testando

Agora, vejamos o desempenho do nosso modelo em dados não vistos:

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))   

Observe que o código é exatamente o mesmo usado para nossos propósitos de validação.

Usando o modelo e treinando apenas por 6 épocas, parecemos ter uma acurácia de aproximadamente 78,8% no conjunto de validação.

Acurácia de Teste

Conclusão

Vamos concluir agora o que fizemos neste artigo:

• Nós começamos entendendo a arquitetura e os diferentes tipos de camadas no modelo AlexNet
• Em seguida, carregamos e pré-processamos o conjunto de dados CIFAR-10 usando torchvision
• Então, usamos PyTorch para construir o nosso modelo AlexNet a partir do zero
• Finalmente, treinamos e testamos o nosso modelo no conjunto de dados CIFAR-10, e o modelo pareceu se sair bem no conjunto de testes com pouco treinamento (6 épocas)

Trabalho Futuro

Este artigo fornece uma Introdução sólida e experiência prática, mas você ganhará ainda mais conhecimento explorando mais a fim de descobrir o que else você pode conseguir.

• Você pode tentar usar diferentes conjuntos de dados. Um desses conjuntos de dados é o CIFAR-100, que é uma extensão do conjunto de dados CIFAR-10 com 100 classes
• Você pode experimentar com diferentes hiperparâmetros e ver a melhor combinação deles para o modelo
• Finalmente, você pode tentar adicionar ou remover camadas do conjunto de dados para ver o impacto delas na capacidade do modelo.