Clasificación de imágenes en Pytorch con Google Colab y Kaggle

En este post, aprenderemos cómo descargar un conjunto de imágenes de Kaggle y a cómo entrenar un modelo de clasificación de imágenes en Pytorch desde Google Colab.

Introducción

En este post, vamos a acceder a un conjunto de datos de imágenes que se encuentra en Kaggle, que para quien no lo conozca, es una plataforma en la que se pueden encontrar diferentes competiciones de problemas de Machine Learning. Crearemos una Red Neuronal Convolucional (Convolutional Neural Network en inglés) o CNN desde cero, la entrenaremos y la intentaremos mejorar.

Hemos llevado a cabo todo el proceso utilizando Google Colab, por lo que si no sabemos cómo instalar librerías de Deep Learning o directamente no tenemos una GPU para agilizar el entrenamiento de modelos, lo recomendamos. Por aquí tenemos un artículo en el que proporcionamos una introducción a esta maravillosa herramienta.

Dataset de Kaggle

Para la clasificación de imágenes con Pytorch hemos optado por seleccionar un conjunto de datos de imágenes de flores de la plataforma de kaggle. Para acceder a ellas hay dos opciones. La primera es descargarlo pulsando directamente en el botón Download.

La segunda opción sería utilizar la API de Kaggle para descargar las imágenes. Esta es la opción que hemos seguido.

Creación de la clave de la API de Kaggle

Para poder utilizar la API de Kaggle, es necesario que creemos una clave. Para ello, debemos pulsar arriba a la derecha, donde está el icono de nuestro perfil y seguidamente en My Account.

A continuación, bajamos hasta que veamos lo siguiente:

Seleccionamos Creare New API Token. Esto generará un archivo llamado kaggle.json. Si estamos usando Google Colab, como recomendabamos arriba, tenemos que subir este archivo a nuestra cuenta de Google Drive.

Descarga de datos con la API de Kaggle

En primer lugar, conectamos nuestra cuenta de Google Drive, con nuestro Notebook de Colab.

from google.colab import drive
drive.mount('/content/drive')

A continuación, copiamos el archivo kaggle.json desde donde lo hayamos subido a la ruta /root. Nosotros lo hemos subido a la carpeta ra?z de Google Drive, pero si lo hemos puesto dentro de alguna carpeta, deberemos indicar la ruta completa.

!chmod 600 drive/'My Drive'/kaggle.json
!mkdir /root/.kaggle/
!cp drive/'My Drive'/kaggle.json /root/.kaggle/

Descargamos el conjunto de datos como un .zip y descomprimimos el archivo.

!kaggle datasets download -d alxmamaev/flowers-recognition
!unzip flowers-recognition.zip -d ./data/

Si entramos en la carpeta data que acabamos de crear, veremos que se han creado 6 directorios nuevos. Cinco se corresponden con las clases que queremos clasificar, pero hay un sexto llamado flowers, que es el mismo conjunto de datos de nuevo, debe de haber sido un error de los autores al subir el dataset, así que eliminaremos esta carpeta.

!rm -rf data/flowers/flowers

Analizando las categorías

Primero cargaremos un conjunto de librerías que nos harán falta a lo largo del proyecto.

import torch
from torchvision.datasets import ImageFolder
import matplotlib.pyplot as plt
from torchvision import transforms
import torchvision
import numpy as np
from torch.utils.data import DataLoader

A continuación cargamos las imágenes para poder analizarlas.

transform = transforms.Compose([transforms.Resize((224, 224)),
                                transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5),
                                                     (0.5, 0.5, 0.5))])
dataset = ImageFolder('data/flowers', transform=transform)
classes = dataset.classes
data_loader = DataLoader(dataset, batch_size=20, shuffle=True)

Creamos una función auxiliar para mostrar las imágenes y mostramos las 20 primeras.

def imshow(img):
    img = img / 2 + 0.5 
    plt.imshow(np.transpose(img, (1, 2, 0)))
dataiter = iter(data_loader)
images, labels = dataiter.next()
fig = plt.figure(figsize=(20, 4))
for idx in np.arange(20):
    ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])
    imshow(images[idx])
    ax.set_title(classes[labels[idx]])

A continuación analizaremos el número de imágenes por categoría.

labels = {}
for label in classes:
  labels[label] = 0
data_loader = DataLoader(dataset, batch_size=1, shuffle=True)
for data in data_loader:
  img, label = data
  labels[classes[label.item()]] += 1
print(labels)

Podemos observar que nuestro conjunto de datos está formado por 4323 imágenes, de las cuáles:

• 729 son margaritas (Daisy)
• 1052 son dientes de le?n (Dandelion)
• 784 son rosas (Rose)
• 734 son girasoles (Sunflower)
• 984 son tulipanes (Tulip)

División en Entrenamiento, Validación y Test

Un paso muy importante cuando entrenamos modelos de Deep Learning y Machine Learning en general, es tener conjuntos de datos para evaluar el modelo en dos fases, durante el entrenamiento y después del entrenamiento. Para el entrenamiento se hará uso de un conjunto de datos de validación, para comprobar que el modelo no esté haciendo overfitting sobre los datos de entrenamiento. Para la fase posterior, se evaluará el modelo con unos datos que nunca haya visto, datos de test.

from torch.utils.data import random_split
train_set, test_set = random_split(dataset, (int(len(dataset) * 0.7) + 1, int(len(dataset) * 0.3)))
train_set, valid_set = random_split(train_set, (int(len(train_set) * 0.7) + 1, int(len(train_set) * 0.3)))

Clasificación de imágenes en Pytorch – Aplicación de CNN desde cero

Arquitectura del modelo de clasificación de imágenes

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(46656, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 5)
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

Entrenamiento del modelo de clasificación

Definimos las funciones de coste y el optimizador que utilizaremos.

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
cnn = Net().to(device)
criterion = nn.CrossEntropyLoss()
# params = cnn.resnet.fc.parameters()
params = cnn.parameters()
optimizer = optim.Adam(params, lr=0.003)

Creamos los DataLoader para cada conjunto de datos, entrenamiento, validación y testeo.

train_loader = DataLoader(train_set, batch_size=64)
valid_loader = DataLoader(valid_set, batch_size=1)
test_loader = DataLoader(test_set, batch_size=1)

Creamos una función para entrenar el modelo.

def train_model(model, train_loader, valid_loader, criterion, optimizer, device):
  total_step = len(train_loader)
  num_epochs = 10
  for epoch in range(num_epochs):
    train_loss = 0.0
    valid_loss = 0.0
    model.train()
    for i, (img, target) in enumerate(train_loader):
      img = img.to(device)
      target = target.to(device)
      optimizer.zero_grad()
      output = model(img)
      
      loss = criterion(output, target)
      loss.backward()
      optimizer.step()
      train_loss += loss.item() * img.size(0)
    model.eval()
    for data, target in test_loader:
      data, target = data.to(device), target.to(device)
      output = model(data)
      loss = criterion(output, target)
      valid_loss += loss.item() * data.size(0)
    train_loss = train_loss / len(train_loader.sampler)
    valid_loss = valid_loss / len(valid_loader.sampler)
    print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(
          epoch, train_loss, valid_loss))

Entrenamos el modelo.

train_model(cnn, train_loader, valid_loader, criterion, optimizer, device)

Creamos una serie de funciones para evaluar el modelo.

def global_accuracy(model, test_loader):
  correct = 0
  total = 0
  model.to("cpu")
  dataiter = iter(test_loader)
  with torch.no_grad():
    for data in dataiter:
      img, label = data
      output = model(img)
      _, predicted = torch.max(output.data, 1)
      total += label.size(0)
      correct += (predicted == label).sum().item()
  print(f"Accuracy: {100 * correct / total}")
def accuracy_per_class(model, test_loader, classes, device):
  class_correct = list(0. for i in range(5))
  class_total = list(0. for i in range(5))
  cnn.to(device)
  with torch.no_grad():
      for data in test_loader:
          images, labels = data
          images = images.to(device)
          labels = labels.to(device)
          outputs = cnn(images)
          _, predicted = torch.max(outputs, 1)
          c = (predicted == labels).squeeze()
          if(c.item()):
            class_correct[labels.item()] += 1
          class_total[labels.item()] += 1

  for i in range(5):
    print(f"{classes[i]} | Correct: {class_correct[i]} | Total: {class_total[i]}" +
          f" | Accuracy: {class_correct[i] / class_total[i]}")

A continuación evaluamos el modelo.

global_accuracy(cnn, test_loader)
accuracy_per_class(cnn, test_loader, classes, device)

Analizando el accuracy, podemos ver que el modelo no muestra resultados extraordinarios.

Mejorando el modelo con Transfer Learning

Probaremos a crear un nuevo modelo utilizando la técnica de aprendizaje por transferencia, en la que cogeremos un modelo preentrenado y le cambiaremos la última capa para que se adecúe a nuestro problema.

class Classifier(nn.Module):
  def __init__(self):
    super(Classifier, self).__init__()
    self.resnet = models.resnet34(pretrained=True)
    self.resnet.fc = nn.Linear(self.resnet.fc.in_features,
                               5)
  def forward(self, image):
    output = self.resnet(image)
    return output

Entrenamiento del modelo ResNet

cnn = Classifier().to(device)
criterion = nn.CrossEntropyLoss()
params = cnn.resnet.fc.parameters()
optimizer = optim.Adam(params, lr=0.003)
train_model(cnn, train_loader, valid_loader, criterion, optimizer, device)
global_accuracy(cnn, test_loader)
accuracy_per_class(cnn, test_loader, classes, device)

Ahora podemos analizar los resultados de este modelo.

Como podemos observar, este modelo que ha sido entrenado utilizando aprendizaje por transferencia, ha mostrado unos resultados muy superiores al modelo entrenado desde cero.

Guardando modelos

Una vez hemos entrenado un modelo y a su vez, este muestra buenos resultados, seguramente queramos almacenar este modelo para utilizarlo posteriormente.

torch.save(cnn.state_dict(), "cnn.pt")

Cargando modelos

Si queremos cargar un modelo que se ha almacenado previamente, debemos primero definir la clase del modelo que queremos cargar y postetiormente, cargar su estado.

model = Classifier()
model.load_state_dict(torch.load('cnn.pt'))

Conclusiones de la Clasificación de Imógenes con Pytorch

En este artículo hemos visto cómo descargar un conjunto de datos de Kaggle utilizando su API, hemos entrenado una CNN desde cero utilizando Pytorch y hemos visto los grandes beneficios de utilizar aprendizaje por transferencia.

Espero que les haya gustado. Si tienen cualquier duda, no duden en dejarla en los comentarios.