Entrenamiento del modelo clasificador de imágenes con PyTorch

En la fase anterior de este tutorial adquirimos el conjunto de datos que usaremos para entrenar nuestro clasificador de imágenes con PyTorch. Ahora es el momento de usar esos datos.

Para entrenar el clasificador de imágenes con PyTorch, es necesario completar los pasos siguientes:

1. Cargue los datos. Si ya ha realizado el paso anterior de este tutorial, ya tiene completado este paso.
2. Defina una red neuronal de convolución.
3. Defina una función de pérdida.
4. Entrene el modelo según los datos de entrenamiento.
5. Pruebe la red en los datos de prueba.

Defina una red neuronal de convolución.

Para crear una red neuronal con PyTorch, usará el paquete torch.nn. Este paquete contiene módulos, clases extensibles y todos los componentes necesarios para compilar redes neuronales.

Aquí creará una red neuronal de convolución básica (CNN) para clasificar las imágenes del conjunto de datos CIFAR10.

Una CNN es una clase de redes neuronales, definidas como redes neuronales multicapa diseñadas para detectar características complejas en los datos. Se usan frecuentemente en aplicaciones de Vision.

Nuestra red se estructurará mediante las 14 capas siguientes:

Conv -> BatchNorm -> ReLU -> Conv -> BatchNorm -> ReLU -> MaxPool -> Conv -> BatchNorm -> ReLU -> Conv -> BatchNorm -> ReLU -> Linear.

La capa de convolución

La capa de convolución es una capa principal de CNN que le permite detectar características en imágenes. Cada una de las capas tiene un número de canales dedicados a detectar características específicas en imágenes y varios kernels que definen el tamaño de la característica detectada. Por lo tanto, una capa de convolución con 64 canales y un tamaño de kernel de 3 x 3 detectará 64 características distintas, cada una de ellas con un tamaño 3 x 3. Al definir una capa de convolución, se proporciona el número de canales de entrada, el número de canales de salida y el tamaño del kernel. El número de canales de salida de la capa actúa como el número de canales de entrada en la capa siguiente.

Por ejemplo: una capa de convolución con 3 canales de entrada, 10 canales de salida y un tamaño de kernel de 6 obtendrá la imagen RGB (3 canales) como entrada y aplicará 10 detectores de características a las imágenes con un tamaño de kernel de 6 x 6. Los tamaños de kernel más pequeños reducirán el tiempo de cálculo y el uso compartido del peso.

Otras capas

Las siguientes capas también forman parte de la red:

• La capa ReLU es una función de activación que se usa para confirmar que todas las características entrantes tengan el valor 0 o superior. Al aplicar esta capa, cualquier número menor que 0 se cambia a cero, mientras que otros se mantienen igual.
• la capa BatchNorm2d aplica la normalización en las entradas para tener una media de cero y una varianza unitaria, y así aumentar la precisión de la red.
• La capa MaxPool le permite asegurarse de que la ubicación de un objeto en una imagen no afectará a la capacidad de la red neuronal de detectar sus características específicas.
• La capa Linear forma parte de las capas finales de la red, que calculan las puntuaciones de cada una de las clases. En el conjunto de datos CIFAR10, hay diez clases de etiquetas. La etiqueta con la puntuación más alta será la que prediga el modelo. En la capa lineal, debe especificar el número de características de entrada y el número de características de salida que deben corresponder al número de clases.

¿Cómo funciona una red neuronal?

CNN es una red de avance de fuentes. Durante el proceso de entrenamiento, la red procesará la entrada a través de todas las capas, calculará las pérdidas para comprender cuánto se aleja la etiqueta predicha de la imagen de la etiqueta correcta y propagará los degradados de nuevo a la red para actualizar los pesos de las capas. Al realizar una iteración sobre un enorme conjunto de datos de entrada, la red "aprenderá" a establecer lo que pesa para lograr los mejores resultados.

Una función directa calcula el valor de la función de pérdida, y la función de retroceso calcula los degradados de los parámetros que se pueden aprender. Al crear la red neuronal con PyTorch, solo tiene que definir la función de directa. La función de retroceso se definirá automáticamente.

1. Copie el código siguiente en el archivo PyTorchTraining.py de Visual Studio para definir la CCN.

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

# Define a convolution neural network
class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=5, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(12)
        self.conv2 = nn.Conv2d(in_channels=12, out_channels=12, kernel_size=5, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(12)
        self.pool = nn.MaxPool2d(2,2)
        self.conv4 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=5, stride=1, padding=1)
        self.bn4 = nn.BatchNorm2d(24)
        self.conv5 = nn.Conv2d(in_channels=24, out_channels=24, kernel_size=5, stride=1, padding=1)
        self.bn5 = nn.BatchNorm2d(24)
        self.fc1 = nn.Linear(24*10*10, 10)

    def forward(self, input):
        output = F.relu(self.bn1(self.conv1(input)))      
        output = F.relu(self.bn2(self.conv2(output)))     
        output = self.pool(output)                        
        output = F.relu(self.bn4(self.conv4(output)))     
        output = F.relu(self.bn5(self.conv5(output)))     
        output = output.view(-1, 24*10*10)
        output = self.fc1(output)

        return output

# Instantiate a neural network model 
model = Network()

Nota:

¿Le interesa obtener más información sobre la red neuronal con PyTorch? Consulte la documentación de PyTorch.

Defina una función de pérdida.

Una función de pérdida calcula un valor que estima a qué distancia está la salida del destino. El objetivo principal es reducir el valor de la función de pérdida cambiando los valores del vector de peso volviendo a propagar los valores en redes neuronales.

Tenga en cuenta que el valor de pérdida es diferente de la precisión del modelo. La función de pérdida le permite comprender cómo se comporta un modelo después de cada iteración de optimización en el conjunto de entrenamiento. La precisión del modelo se calcula en los datos de prueba y muestra el porcentaje de la predicción correcta.

En PyTorch, el paquete de red neuronal contiene varias funciones de pérdida que forman los bloques de creación de redes neuronales profundas. En este tutorial, usará una función de pérdida de clasificación basada en la función para definir la pérdida con la pérdida de entropía cruzada de clasificación y un optimizador de tipo Adam. La tasa de aprendizaje (lr) establece el control de cuánto está ajustando los pesos de nuestra red con respecto al gradiente de pérdida. Establecerá este valor en 0,001. Cuanto menor sea, más lento será el entrenamiento.

1. Copie el código siguiente en el archivo PyTorchTraining.py de Visual Studio para definir la función de pérdida y un optimizador.

from torch.optim import Adam
 
# Define the loss function with Classification Cross-Entropy loss and an optimizer with Adam optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.001, weight_decay=0.0001)

Entrene el modelo según los datos de entrenamiento.

Para entrenar el modelo, debe recorrer en bucle el iterador de datos, alimentar las entradas a la red y optimizarlas. PyTorch no tiene una biblioteca dedicada para usarla en GPU, pero puede definir manualmente el dispositivo de ejecución. El dispositivo será una GPU de Nvidia si existe en la máquina, o la CPU si no es así.

1. Agregue el siguiente código al archivo PyTorchTraining.py.

from torch.autograd import Variable

# Function to save the model
def saveModel():
    path = "./myFirstModel.pth"
    torch.save(model.state_dict(), path)

# Function to test the model with the test dataset and print the accuracy for the test images
def testAccuracy():
    
    model.eval()
    accuracy = 0.0
    total = 0.0
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            # run the model on the test set to predict labels
            outputs = model(images.to(device))
            # the label with the highest energy will be our prediction
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            accuracy += (predicted == labels.to(device)).sum().item()
    
    # compute the accuracy over all test images
    accuracy = (100 * accuracy / total)
    return(accuracy)


# Training function. We simply have to loop over our data iterator and feed the inputs to the network and optimize.
def train(num_epochs):
    
    best_accuracy = 0.0

    # Define your execution device
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("The model will be running on", device, "device")
    # Convert model parameters and buffers to CPU or Cuda
    model.to(device)

    for epoch in range(num_epochs):  # loop over the dataset multiple times
        running_loss = 0.0
        running_acc = 0.0

        for i, (images, labels) in enumerate(train_loader, 0):
            
            # get the inputs
            images = Variable(images.to(device))
            labels = Variable(labels.to(device))

            # zero the parameter gradients
            optimizer.zero_grad()
            # predict classes using images from the training set
            outputs = model(images)
            # compute the loss based on model output and real labels
            loss = loss_fn(outputs, labels)
            # backpropagate the loss
            loss.backward()
            # adjust parameters based on the calculated gradients
            optimizer.step()

            # Let's print statistics for every 1,000 images
            running_loss += loss.item()     # extract the loss value
            if i % 1000 == 999:    
                # print every 1000 (twice per epoch) 
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 1000))
                # zero the loss
                running_loss = 0.0

        # Compute and print the average accuracy fo this epoch when tested over all 10000 test images
        accuracy = testAccuracy()
        print('For epoch', epoch+1,'the test accuracy over the whole test set is %d %%' % (accuracy))
        
        # we want to save the model if the accuracy is the best
        if accuracy > best_accuracy:
            saveModel()
            best_accuracy = accuracy

Pruebe el modelo con los datos de prueba.

A continuación, pruebe el modelo con el lote de imágenes del conjunto de pruebas.

1. Agregue el siguiente código al archivo PyTorchTraining.py.

import matplotlib.pyplot as plt
import numpy as np

# Function to show the images
def imageshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


# Function to test the model with a batch of images and show the labels predictions
def testBatch():
    # get batch of images from the test DataLoader  
    images, labels = next(iter(test_loader))

    # show all images as one image grid
    imageshow(torchvision.utils.make_grid(images))
   
    # Show the real labels on the screen 
    print('Real labels: ', ' '.join('%5s' % classes[labels[j]] 
                               for j in range(batch_size)))
  
    # Let's see what if the model identifiers the  labels of those example
    outputs = model(images)
    
    # We got the probability for every 10 labels. The highest (max) probability should be correct label
    _, predicted = torch.max(outputs, 1)
    
    # Let's show the predicted labels on the screen to compare with the real ones
    print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] 
                              for j in range(batch_size)))

Por último, agregue el código principal. Esto iniciará el entrenamiento del modelo, guardará el modelo y mostrará los resultados en la pantalla. Tenga en cuenta que solo ejecutará dos iteraciones [train(2)] en el conjunto de entrenamiento, por lo que el proceso de entrenamiento no llevará demasiado tiempo.

2. Agregue el siguiente código al archivo PyTorchTraining.py.

if __name__ == "__main__":
    
    # Let's build our model
    train(5)
    print('Finished Training')

    # Test which classes performed well
    testAccuracy()
    
    # Let's load the model we just created and test the accuracy per label
    model = Network()
    path = "myFirstModel.pth"
    model.load_state_dict(torch.load(path))

    # Test with batch of images
    testBatch()

Ejecute la prueba. Asegúrese de que los menús desplegables de la barra de herramientas superior están establecidos en Depurar. Cambie la plataforma de soluciones a x64 para ejecutar el proyecto en la máquina local si el dispositivo es de 64 bits o a x86 si el dispositivo es de 32 bits.

La elección del número de época (esto es, el número de pases completos a través del conjunto de datos de entrenamiento) igual a dos ([train(2)]) dará lugar a la iteración dos veces por todo el conjunto de datos de prueba de 10 000 imágenes. El entrenamiento en la CPU Intel de 8ª generación llevará unos 20 minutos, por lo que el modelo debería lograr más o menos el 65 % de la tasa de éxito en la clasificación de diez etiquetas.

3. Para ejecutar el proyecto, haga clic en el botón Iniciar depuración en la barra de herramientas o pulse F5.

Aparecerá la ventana de la consola y podrá ver el proceso de entrenamiento.

Según lo haya definido, el valor de pérdida se imprimirá cada 1000 lotes de imágenes o cinco veces por cada iteración en el conjunto de entrenamiento. Se espera que el valor de pérdida disminuya con cada bucle.

También verá la precisión del modelo después de cada iteración. Recuerde que la precisión del modelo es diferente del valor de pérdida. La función de pérdida le permite comprender cómo se comporta un modelo después de cada iteración de optimización en el conjunto de entrenamiento. La precisión del modelo se calcula en los datos de prueba y muestra el porcentaje de la predicción correcta. En nuestro caso, nos dirá cuántas imágenes del conjunto de pruebas de 10 000 imágenes pudo clasificar el modelo correctamente después de cada iteración de entrenamiento.

Una vez completado el entrenamiento, verá un resultado similar al siguiente. Los números no serán exactamente iguales (la diferencia depende de muchos factores y no siempre devolverá resultados que se puedan identificar), pero deben ser similares.

Después de ejecutar solo 5 épocas, la tasa de éxito del modelo es del 70 %. Este es un buen resultado para un modelo básico entrenado durante un breve período de tiempo.

Al probar con el lote de imágenes, el modelo obtuvo 7 imágenes correctas de un lote de 10. Esto no está nada mal y es coherente con la tasa de éxito del modelo.

Puede comprobar qué clases puede predecir mejor el modelo. Sencillamente, agregue el código siguiente:

4. Opcional: agregue la siguiente función testClassess al archivo PyTorchTraining.py y agregue una llamada a esta función - testClassess() en la función principal - __name__ == "__main__".

# Function to test what classes performed well
def testClassess():
    class_correct = list(0. for i in range(number_of_labels))
    class_total = list(0. for i in range(number_of_labels))
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            c = (predicted == labels).squeeze()
            for i in range(batch_size):
                label = labels[i]
                class_correct[label] += c[i].item()
                class_total[label] += 1

    for i in range(number_of_labels):
        print('Accuracy of %5s : %2d %%' % (
            classes[i], 100 * class_correct[i] / class_total[i]))

La salida es como sigue:

Pasos siguientes

Ahora que ya tiene un modelo de clasificación, el siguiente paso es convertir el modelo al formato ONNX.