train.py

# %%
import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm

from data_loader import get_train_test_datasets
from model import Model

# %%

# --- Config ---
BATCH_SIZE = 32
NUM_EPOCHS = 100
LR = 1e-3
CSV = "data/labels.csv"
IMG_DIR = "data/img"

# Datasets & loaders
train_dataset, test_dataset = get_train_test_datasets(
    csv_file="data/labels.csv", img_dir="data/img", transform=transforms.ToTensor()
)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# Train CNN
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")

nn = Model().to(device)
optimizer = torch.optim.Adam(nn.parameters(), lr=LR)
criterion = torch.nn.CrossEntropyLoss()

loss_list = []

# %%
# Train loop
for epoch in tqdm(range(NUM_EPOCHS), desc="Epochs"):
    loss_list_epoch = []
    # Itero sobre todos los batches del dataset
    for x, y in train_loader:
        # Movemos los tensores a memoria de GPU
        x = x.to(device)
        y = y.to(device)
        # y = y.squeeze().long()

        # Seteo en cero los gradientes de los parámetros a optimizar
        optimizer.zero_grad()

        # Realizo la pasada forward por la red
        logits = nn(x)
        # print(type(logits), type(y))
        loss = criterion(logits, y)

        # Realizo la pasada backward por la red
        loss.backward()

        # Actualizo los pesos de la red con el optimizador
        optimizer.step()

        # Me guardo el valor actual de la función de pérdida para luego graficarlo
        loss_list.append(loss.data.item())
        loss_list_epoch.append(loss.data.item())

    # Muestro el valor de la función de pérdida cada 100 iteraciones
    if epoch > 0 and epoch % 10 == 0:
        print(
            "Epoch %d, Avg train loss during epoch = %g"
            % (epoch, np.array(loss_list_epoch).mean())
        )

# Muestro la lista que contiene los valores de la función de pérdida
# y una versión suavizada (rojo) para observar la tendencia
plt.figure()
loss_np_array = np.array(loss_list)
plt.plot(loss_np_array, alpha=0.3)
N = 60
running_avg_loss = np.convolve(loss_np_array, np.ones((N,)) / N, mode="valid")
plt.plot(running_avg_loss, color="red")
plt.title("Función de pérdida durante el entrenamiento")

torch.save(nn.state_dict(), "./cnn.pth")
# %%
"""
EJEMPLO DE COMO PREDECIR

from PIL import Image

cross = Image.open("base_img/equal.png").convert("RGB")
tensor_cross = transforms.ToTensor()(cross)

print(torch.argmax(nn(tensor_cross.unsqueeze(0))))
"""
# %%


def evaluate_model(model, test_loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for x, y in test_loader:
            x = x.to(device)
            y = y.to(device)
            outputs = model(x)
            _, predicted = torch.max(outputs, 1)
            total += y.size(0)
            correct += (predicted == y).sum().item()
    accuracy = correct / total if total > 0 else 0
    print(f"Test Accuracy: {accuracy:.4f}")


evaluate_model(nn, test_loader, device)
# %%