Escrevendo AlexNet derazo em PyTorch

Introdução

Este post é uma continuação da série sobre a construção das redes neurais convolucionais mais populares a partir do zero em PyTorch. Você pode ver o post anterior aqui, onde construímos LeNet5. Neste post, vamos construir AlexNet, uma das Breakthrough Algorithms mais importantes em computação visual.

Vamos começar 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é-processadas aos dados. Em seguida, construiremos o nosso AlexNet a partir do zero usando 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 (variantes 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 de CNNs é essencial. Isso inclui o conhecimento de camadas convolucionais, camadas de pooling, e seu papel na extração de recursos 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 detalhadamente sobre o modelo no artigo de pesquisa original aqui.

Vamos passar por as conclusões chave 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 ao subamostrar. 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 processando 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 totalmente 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 pesados. 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 notebook 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 torchvision (uma biblioteca de apoio para tarefas de visão computacional), carregaremos nosso conjunto 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 o conjunto 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 os conjuntos de treinamento/validação e teste respectivamente
• Começamos definindo a variável normalize com a média e as desvios padrão de cada canal (vermelho, verde e azul) no conjunto de dados. Esses valores 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, para tornar o treinamento 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 sub-selecionamos aleatoriamente a partir 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 no tipo de rótulo que tem. Isto aumentará a eficácia de nosso 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 é considerado uma boa prática em geral. Os 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 seu 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

Definição do 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) no PyTorch é definir uma classe que herda de nn.Module pois contém muitos dos métodos que vamos precisar utilizar
• Existem dois passos principais depois disso. O primeiro é inicializar as camadas que vamos usar em nossa CNN dentro de __init__, e o outro é 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 pooling máximo usando a função nn.MaxPool2D. O que é legal em PyTorch é que podemos combinar a camada convolucional, função de ativação e pooling máximo em uma única camada (eles serão aplicados 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), dropout (nn.Dropout) juntamente com a função de ativação ReLu (nn.ReLU) e combinando estes 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 da treinamento, precisamos configurar alguns hiperparâmetros, como a função de perda e o otimizador a serem usados, juntamente com o tamanho do lote, taxa de aprendizado 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 simples hiperparâmetros (epocas, tamanho de lote e taxa de aprendizagem) e inicializando o nosso modelo com 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). Depois definimos nossa função de custo como perda de cross-entropia e o otimizador como Adam. Existem muitas escolhas para estes, mas essas tendem a dar bons resultados com o modelo e os dados fornecidos. Finalmente, definimos total_step para melhor rastrear os passos durante o treinamento

Treino

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)

        # Adiante e otimizar
        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)) 

Vejamos o que o código faz:

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

A saída pode ser vista da seguinte forma:

Perda de Treinamento e Precisão de Validação

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

Teste

Agora, vemos como o nosso modelo performa 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))   

Note que o código é exatamente o mesmo usado para nossos fins de validação.

Usando o modelo e treinando apenas por 6 épocas, obtemos cerca de 78,8% de precisão no conjunto de validação.

Precisão de Teste

Conclusão

Vamos concluir agora o que fizemos neste artigo:

• Começamos entendendo a arquitetura e os diferentes tipos de camadas do 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 de raiz
• Finalmente, treinamos e testamos o nosso modelo no conjunto de dados CIFAR-10, e o modelo pareceu performar bem no conjunto de testes com pouco treinamento (6 épocas)

Trabalho Futuro

Este artigo fornece uma boa introdução e experiência prática, mas você vai ganhar ainda mais conhecimento explorando mais a fundo e descobrindo o que mais você pode conseguir.

• Você pode tentar usar diferentes conjuntos de dados. Um desses conjuntos é 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.