En resumen, en este pequeño tutorial aprenderemos a usar PyTorch para crear, entrenar y predecir con nuestra primera red neuronal.

Primeros pasos

Lo primero que debemos hacer es importar torch y algunos módulos de esta librería:

• torch.nn: La librería de redes neuronales que utilizaremos para crear nuestro modelo.
• torch.autograd: En concreto el módulo Variable de esta librería que se encarga de manejar las operaciones de los tensores.
• torchvision.datasets: El módulo que ayudará a cargar el conjunto de datos que vamos a utilizar y explicaremos más adelante.
• torchvision.transforms: Este módulo contiene una serie de funciones que nos ayudarán modificando el dataset.
• torch.optim: De aquí usaremos el optimizador para entrenar la red neuronal y modificar sus pesos.

Para ello ejecutamos el siguiente código:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torch.optim as optim

Dataset a utilizar

Una vez importados los módulos, debemos cargar el conjunto de datos que vamos a utilizar, en este tutorial vamos a utilizar el conjunto conocido como MNIST. Se trata de un conjunto de imágenes de dígitos escritos a mano, contiene un total de 60000 imágenes para entrenamiento y 10000 para test. Todos los dígitos están normalizados en tamaño y centrados en la imagen de tamaño 28×28 en escala de grises. El objetivo de esta base de datos es clasificar cada imagen diciendo a que número entre el 0 y el 9 pertenece.

Ejemplo de los dígitos de MNIST (fuente)

Para cargar de forma rápida el dataset vamos a utilizar el modulo datasets de torchvision, además debemos definir que transformaciones vamos a aplicarle a todas las muestras. En nuestro caso únicamente vamos a pasar el objeto obtenido a Tensor con lo que PyTorch podrá realizar cálculos, torchvision.transforms contiene una gran cantidad de transformaciones para aplicar a los datasets. Es recomendable normalizar nuestro dataset, este dataset viene por defecto normalizado valores entre 0 y 1. También vamos a fijar la semilla del generador aleatorio para poder replicar los experimentos.

torch.manual_seed(123) #fijamos la semilla
trans = transforms.Compose([transforms.ToTensor()]) #Transformador para el dataset

Después definimos donde vamos a descargar los datos, en nuestro caso en la carpeta data dentro de la que estamos utilizando. La herramienta datasets de torchvision es muy fácil y cómoda para usar, simplemente debemos elegir el dataset entre los que tiene disponibles, debemos indicar la ruta a la carpeta con el parámetro root, si queremos el dataset de entrenamiento (en caso contrario obtendremos el de test), las transformaciones que se le van a aplicar que hemos definido anteriormente y si queremos descargar el dataset en el caso de que no lo tengamos almacenado ya. Para ello utilizaremos el siguiente código:

root="./data"
train_set = dset.MNIST(root=root, train=True, transform=trans, download=True)
test_set = dset.MNIST(root=root, train=False, transform=trans)

Lo siguiente que debemos hacer es definir el tamaño de batch, que es simplemente el tamaño de conjunto de datos que vamos a pasar a la red a la vez, es decir, en vez de clasificar y actualizar los pesos muestra por muestra vamos a hacerlo cada conjunto de x muestras, esto ayuda a paralelizar la red pero debemos tener cuidado porque tamaños muy grandes de batch pueden requerir grandes cantidades de memoria RAM, ya sea en la GPU o en la CPU. Después definimos un Dataloader para el dataset que hemos descargado. Un Dataloader no es más que un generador que nos proporcionará las muestras en grupos de un tamaño dado, además contiene ciertas funciones de gran ayuda como la de reordenar aleatoriamente el dataset en cada iteración (shuffle). El código para generar los dataloader de entrenamiento y test es el siguiente:

batch_size = 128
train_loader = torch.utils.data.DataLoader(
                 dataset=train_set,
                 batch_size=batch_size,
                 shuffle=True)
test_loader = torch.utils.data.DataLoader(
                dataset=test_set,
                batch_size=batch_size,
                shuffle=False)

Podemos saber la cantidad de conjuntos (batches) de tamaño 128 tenemos simplemente utilizando el método len de Python. En el caso de que el dataset no sea un múltiplo del tamaño de batch, el último batch será de un tamaño reducido.

print ('Trainning batch number: {}'.format(len(train_loader)))
print ('Testing batch number: {}'.format(len(test_loader)))

Creación del modelo y entrenamiento

Una vez tenemos el dataset debemos decidir que topología de red neuronal vamos a utilizar, en nuestro caso vamos a optar por una muy simple con una sola capa oculta que contiene 256 neuronas. Para ello vamos a definir una clase que desciende de nn.Module. En el __init__ definiremos las capas que va a tener la red y la función de pérdida que vamos a utilizar. La primera capa será la capa oculta y recibirá el tensor de entrada con los datos que tienen un tamaño de 28×28, por tanto la capa tendrá un tamaño de entrada de 28*28 y un tamaño de salida de 256. La segunda capa será la capa de salida que recibirá los datos de la capa oculta con un tamaño de 256 y su salida será igual al número de clases, en nuestro caso 10 clases. La función de pérdida que hemos escogido para nuestra red es la entropía cruzada.

Después debemos definir la función forward, es decir, las operaciones que se van a realizar desde la entrada de la red hasta su salida. Aquí vamos a definir las funciones que se aplicarán de forma sucesiva a la red, en nuestro caso es muy sencillo. Primero debemos transformar la entrada que será una matriz de 28×28 en un vector de 784, para ello utilizaremos la función view, equivalente a reshape en numpy (el -1 se hace para que la primera dimensión, el tamaño de batch no sea alterada). Después aplicamos un función de activación a la salida de la capa oculta, en este caso hemos optado por la función rectified linear unit. Por último aplicamos la función softmax al tensor devuelto por la capa de salida, el parámetro dim indica en que dimensión se efectuará la función. Por último calculamos la función de pérdida, no es una práctica recomendable añadir el cálculo de dicha función al método forward ya que en el caso de únicamente querer predecir sería un cálculo absurdo pero en este caso lo haremos para que nos sea más sencillo entrenar la red. El código de dicha clase será el siguiente:

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(28*28, 256)#capa oculta
        self.fc2 = nn.Linear(256, 10)#capa de salida
        self.loss_criterion = nn.CrossEntropyLoss()#Función de pérdida
        
    def forward(self, x, target):
        x = x.view(-1, 28*28)#transforma las imágenes de tamaño (n, 28, 28) a (n, 784)
        x = F.relu(self.fc1(x))#Función de activación relu en la salida de la capa oculta
        x = F.softmax(self.fc2(x), dim=1)#Función de activación softmax en la salida de la capa oculta
        loss = self.loss_criterion(x, target)#Calculo de la función de pérdida
        return x, loss

Ahora que tenemos el modelo de red neuronal creamos un objeto de dicha clase y el optimizador que vamos a utilizar para ajustar los pesos de la red. En nuestro caso será el descenso por gradiente estocástico el optimizador que utilizaremos con un ratio de aprendizaje de 0.1 y un momento de 0.9.

model = MLP()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

Entrenamiento del modelo

Ahora vamos a desarrollar una función para poder entrenar la red, ya que hay que hacer el forward tanto para entrenamiento como para test vamos a utilizar la misma función. Dicha función recibirá el modelo, el dataloader, el optimizador y si estamos entrenando o no. Debemos tener un bucle que itere sobre cada batch del dataloader, en cada iteración el dataloader devuelve una tupla que contiene los datos de entrada y la salida esperada. Para cada tupla debemos poner a cero los gradientes de los pesos (si estamos entrenando). Después convertimos la entrada y el objetivo a un objeto llamado Variable dentro de autograd, esto hace que el objeto sea capaz de almacenar los gradientes y propagarlos. Una vez hecho esto podemos realizar el forward sobre la red y obtendremos la predicción (score) y el valor que devuelve la función de pérdida (loss). Con el valor de la predicción calculamos la clase predicha simplemente eligiendo la que obtenga el mayor valor y después comprobamos cuantas han sido correctas. En el caso de estar entrenando devemos propagar los gradientes con el algoritmo backward y después utilizar el optimizador para actualizar los pesos de la red. La función que acabamos de describir sería la siguiente:

def evaluate(model, dataset_loader, optimizer, train=False):
    correct_cnt, ave_loss = 0, 0#Contador de aciertos y acumulador de la función de pérdida
    count = 0#Contador de muestras
    for batch_idx, (x, target) in enumerate(dataset_loader):
        count += len(x)#sumamos el tamaño de batch, esto es porque n_batches*tamaño_batch != n_muestras
        if train:
            optimizer.zero_grad()#iniciamos a 0 los valores de los gradiente
        x, target = Variable(x), Variable(target)#Convertimos el tensor a variable del modulo autograd
        score, loss = model(x, target)#realizamos el forward
        _, pred_label = torch.max(score.data, 1)#pasamos de one hot a número
        correct_cnt += (pred_label == target.data).sum()#calculamos el número de etiquetas correctas
        ave_loss += loss.data[0]#sumamos el resultado de la función de pérdida para mostrar después
        if train:
            loss.backward()#Calcula los gradientes y los propaga 
            optimizer.step()#adaptamos los pesos con los gradientes propagados
    accuracy = correct_cnt/count#Calculamos la precisión total
    ave_loss /= count#Calculamos la pérdida media
    print ('==>>>loss: {:.6f}, accuracy: {:.4f}'.format(ave_loss, accuracy))#Mostramos resultados

Entrenar nuestro modelo ahora es muy sencillo, por ejemplo, si queremos realizar el proceso de entrenamiento durante 10 epochs testeando en cada una simplemente ejecutamos el siguiente código:

for epoch in range(10):
    print("Epoch: {}".format(epoch))
    print("Train")
    evaluate(model, train_loader, optimizer, train=True)
    print("Test")
    evaluate(model, test_loader, optimizer, train=False)

Guardar y cargar el modelo entrenado

Si hemos seguido todos los pasos, incluido fijar la semilla de PyTorch obtendremos una precisión de Test del 92.73 % expresado en tanto por uno.
Una vez entrenado el modelo podemos guardarlo y cargarlo fácilmente, para ello tenemos dos opciones, guardar el modelo entero o guardar solo los pesos. La primera opción es quizá la más sencilla:

torch.save(model, "nombre_del_modelo")
model = torch.load("nombre_del_modelo")#Carga el modelo

Una vez cargado, el modelo es el mismo al que teníamos antes y podemos volver a entrenar o pasar el test. En el caso de querer guardar solo los pesos, la operación es igual de sencilla aunque a la hora de cargarlos debemos generar primero el modelo al igual que hemos hecho antes de entrenarlo y después cargar los pesos:

torch.save(model.state_dict(),"nombre_del_modelo")
model = MLP()#Generamos un nuevo modelo
model.load_state_dict(torch.load("nombre_del_modelo"))#Cargamos los pesos anteriores

Este método es el recomendado para realizar el guardado y así evitar incompatibilidades debidas a la serialización del objeto.

Modelo completo con integración en CUDA

Para hacer más sencilla la integración hemos diseñado la siguiente clase que permite realizar el entrenamiento y clasificación de una forma más sencilla y permite su integración con CUDA para el cálculo en GPU. De esta forma tendremos una clase base donde lo único que modificaremos será el __init__ y el forward de la clase NET para entrenar una topología diferente.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torch.optim as optim
import time

CUDA = torch.cuda.is_available()

class NET(nn.Module):
    def __init__(self):
        super(NET, self).__init__()
        self.fc1 = nn.Linear(28*28, 256)#capa oculta
        self.fc2 = nn.Linear(256, 10)#capa de salida
        self.loss_criterion = nn.CrossEntropyLoss()#Función de pérdida
        
    def forward(self, x, target):
        x = x.view(-1, 28*28)#transforma las imágenes de tamaño (n, 28, 28) a (n, 784)
        x = F.relu(self.fc1(x))#Función de activación relu en la salida de la capa oculta
        x = F.softmax(self.fc2(x), dim=1)#Función de activación softmax en la salida de la capa oculta
        return x
        

class NN():
    def __init__(self):
        self.model = NET()
        if CUDA:
            self.model.cuda()
        self.optimizer = optim.SGD(self.model.parameters(), lr=0.01, momentum=0.9)

    def epoch_step(self, dataset_loader, train=False, verbose=2):
        correct_cnt, ave_loss = 0, 0#Contador de aciertos y acumulador de la función de pérdida
        count = 0#Contador de muestras
        for batch_idx, (x, target) in enumerate(dataset_loader):
            start_time_epch = time.time()
            count += len(x)#sumamos el tamaño de batch, esto es porque n_batches*tamaño_batch != n_muestras
            if train:
                self.optimizer.zero_grad()#iniciamos a 0 los valores de los gradiente
            x, target = Variable(x), Variable(target)#Convertimos el tensor a variable del modulo autograd
            if CUDA:
                x = x.cuda()
                target = target.cuda()
            score = self.model(x, target)#realizamos el forward
            loss = self.model.loss_criterion(score, target)
            _, pred_label = torch.max(score.data, 1)#pasamos de one hot a número
            correct_cnt_epch = (pred_label == target.data).sum()#calculamos el número de etiquetas correctas
            correct_cnt += correct_cnt_epch
            ave_loss += loss.data[0]#sumamos el resultado de la función de pérdida para mostrar después
            if train:
                loss.backward()#Calcula los gradientes y los propaga 
                self.optimizer.step()#adaptamos los pesos con los gradientes propagados
            elapsed_time_epch = time.time() - start_time_epch

            if verbose == 1:
                elapsed_time_epch = time.strftime("%Hh,%Mm,%Ss", time.gmtime(elapsed_time_epch))
                print ('\t\tbatch: {} loss: {:.6f}, accuracy: {:.4f}, time: {}'.format(
                    batch_idx, loss.data[0], correct_cnt_epch/len(x), elapsed_time_epch))

        accuracy = correct_cnt/count#Calculamos la precisión total
        ave_loss /= count#Calculamos la pérdida media

        return ave_loss, accuracy

    def train(self, epoch, train_loader, test_loader=None, verbose=2):
        for epoch in range(epoch):
            if verbose > 0:
                print("\n***Epoch {}***\n".format(epoch))
                if verbose == 1:
                    print("\tTraining:")
            start_time = time.time()
            ave_loss, accuracy = self.epoch_step(train_loader, train=True, verbose=verbose)
            elapsed_time = time.time() - start_time
            if verbose == 2:
                elapsed_time = time.strftime("%Hh,%Mm,%Ss", time.gmtime(elapsed_time))
                print ('\tTraining loss: {:.6f}, accuracy: {:.4f}, time: {}'.format(
                    ave_loss, accuracy, elapsed_time))

            if test_loader != None:
                start_time = time.time()
                ave_loss, accuracy = self.epoch_step(test_loader, train=False, verbose=verbose)
                elapsed_time = time.time() - start_time
                if verbose == 2:
                    elapsed_time = time.strftime("%Hh,%Mm,%Ss", time.gmtime(elapsed_time))
                    print ('\tTest loss: {:.6f}, accuracy: {:.4f}, time: {}'.format(
                        ave_loss, accuracy, elapsed_time))

    def evaluate(self, test_loader, verbose=2):
        print("\n***Evaluate***\n")
        start_time = time.time()
        ave_loss, accuracy = self.epoch_step(test_loader, train=False, verbose=verbose)
        elapsed_time = time.time() - start_time
        if verbose == 2:
            elapsed_time = time.strftime("%Hh,%Mm,%Ss", time.gmtime(elapsed_time))
            print ('\tloss: {:.6f}, accuracy: {:.4f}, time: {}'.format(
                ave_loss, accuracy, elapsed_time))

    def save_weights(self, path):
        torch.save(self.model.state_dict(), path)
    
    def load_weights(self, path):
        self.model.load_state_dict(torch.load(path))


if __name__ == "__main__":
    torch.manual_seed(123) #fijarmos la semilla
    epochs = 10

    trans = transforms.Compose([transforms.ToTensor()]) #Transformador para el dataset
    root="./data"
    train_set = dset.MNIST(root=root, train=True, transform=trans, download=True)
    test_set = dset.MNIST(root=root, train=False, transform=trans)

    batch_size = 128
    train_loader = torch.utils.data.DataLoader(
                    dataset=train_set,
                    batch_size=batch_size,
                    shuffle=True,
                    pin_memory=CUDA)
    test_loader = torch.utils.data.DataLoader(
                    dataset=test_set,
                    batch_size=batch_size,
                    shuffle=False,
                    pin_memory=CUDA)

    print ('Trainning batch number: {}'.format(len(train_loader)))
    print ('Testing batch number: {}'.format(len(test_loader)))

    net = NN()
    net.train(epochs, train_loader, test_loader=test_loader, verbose=2)
    net.evaluate(test_loader, verbose=2)
    net.save_weights("./weights")

En el inicio tenemos una variable booleana llamada CUDA, esta se pondrá a True si disponemos de una GPU NVIDIA en nuestro ordenador. Pese a que tengamos una GPU NVIDIA, no todas puede ejecutar PyTorch ya que necesita que la GPU tenga una capacidad de computo 3.0 por lo menos, podemos consultar una tabla que nos indica que capacidad tiene la nuestra aquí. En el caso de que la gráfica tenga una potencia inferior, posiblemente recibiremos varios errores donde uno de ellos será el siguiente:

torch.save(model.state_dict(),"nombre_del_modelo")
model = MLP()#Generamos un nuevo modelo
model.load_state_dict(torch.load("nombre_del_modelo"))#Cargamos los pesos anteriores

Para solucionar el error debemos ejecutar el código en la CPU, para ello simplemente ponemos la variable CUDA a False y ejecutará todos los cálculos en nuestra CPU.