Training a Vision Model (I): First Steps with a CNN (Convolutional neural network)

This is the first installment of a practical series on how to train a computer vision model from scratch. We don't start from abstract theory: we start with a real, working script, and we take it apart piece by piece to understand what each line does and why.

In this first article, we set up the skeleton: the folder structure, the training script, and the architecture of the convolutional neural network (CNN) that we will use as a base.

1. Starting Point: Project Structure

Before writing code, this is how we organize the project directory:

testENTRENAR_IA/
├── data/
│   └── cifar-10/     # images and labels
├── models/
│   └── modelo_cifar10.pth     # trained weights
└── train_cifar10.py            # training script

Nothing exotic: a folder for input data, another where already trained weights are saved, and the main script. This structure is repeated in almost any training project, so it is good to get used to it from day one.

Download: You can download the training dataset from cave.cs.toronto.edu

2. The Training Script, Section by Section

Paths and Device

# os.path.abspath(__file__) gives the full path of this script.
# os.path.dirname() gets only the folder containing it.
# Thus ROOT_DIR always points to "testENTRENAR_IA/", no matter from where you run the script.
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
# We build the rest of the paths from ROOT_DIR, instead of writing them manually,
# so the project works the same on your PC, server, etc.
DATA_DIR = os.path.join(ROOT_DIR, "data", "cifar10")   # folder with the images
TRAIN_DIR = os.path.join(DATA_DIR, "train")             # training subfolder
TEST_DIR = os.path.join(DATA_DIR, "test")               # evaluation subfolder
MODEL_DIR = os.path.join(ROOT_DIR, "models")            # here we will save the already trained model
# Creates the "models/" folder if it doesn't exist yet.
# exist_ok=True prevents the script from failing if the folder is already created.
os.makedirs(MODEL_DIR, exist_ok=True)
# We choose where to train: GPU if one is available (much faster),
# and if not, we automatically fall back to CPU without breaking the script.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

All of this is preparation: where the data is, where the trained model will be saved, and which device it will be trained on. If a GPU is available (`cuda`), it is used; if not, it automatically falls back to CPU.

Transforms and Data Loading (Dataset Review)

# transforms.Compose chains several transformations applied to EACH image,
# in order, before entering the network.
transform = transforms.Compose([
    transforms.Resize((32, 32)),   # all images must be the same size; 32x32 = CIFAR-10 size
    transforms.ToTensor(),          # converts the image (pixels 0-255) to a PyTorch tensor (0.0-1.0)
    transforms.Normalize((0.5,), (0.5,))   # rescales values to a range centered at 0, helping the training converge better
])
# ImageFolder assumes that each subfolder inside TRAIN_DIR is a different class.
# Example: train/hen/, train/duck/, train/goose/ -> 3 classes detected automatically.
train_dataset = ImageFolder(root=TRAIN_DIR, transform=transform)
# DataLoader is responsible for serving images in "batches" instead of one by one.
# batch_size=32 -> we process 32 images at once before updating the model.
# shuffle=True -> shuffles the order in each epoch so the model doesn't memorize the data order.
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

Images do not enter the network "as is": they are resized to a fixed size (32×32, the same used by the CIFAR-10 dataset) and normalized. `ImageFolder` is a PyTorch utility that assumes each subfolder inside `train/` is a distinct class — so if you have `train/hen/`, `train/duck/`, and `train/goose/`, PyTorch automatically detects those three classes without you having to label them manually.

3. The Architecture: The `Net` Class

Here is the heart of the model, represented by the network architecture:

# Every network in PyTorch inherits from nn.Module. This gives it things like
# moving the model to GPU, saving/loading weights, etc., for free.
class Net(nn.Module):
    # __init__ is where we DECLARE the layers. No data passes through here yet;
    # we are just placing the "pieces" ready on the table.
    def __init__(self, num_classes):
        super(Net, self).__init__()   # mandatory: initializes the internal part of nn.Module
        # First convolutional layer:
        # 3  -> input channels (color image: Red, Green, Blue)
        # 16 -> number of filters we will learn (16 different "detectors")
        # 3  -> filter size, 3x3 pixels
        # padding=1 -> adds a 1-pixel border so the image doesn't shrink when convolving
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        # Pooling: reduces the image size by half.
        # (2, 2) -> 2x2 pixel window, keeps the highest value of each window
        self.pool = nn.MaxPool2d(2, 2)
        # Second convolutional layer:
        # 16 -> enters with the 16 filters generated by conv1
        # 32 -> outputs 32 new filters, looking for more complex patterns than the first layer
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        # Dense (fully connected) layer: connects ALL input neurons with ALL output neurons.
        # 32 * 8 * 8 -> size of the flattened tensor coming out of the convolutions (see explanation below)
        # 128 -> size of the "summary" vector generated for each image
        self.fc1 = nn.Linear(32 * 8 * 8, 128)
        # Last layer: one output neuron for each possible class (hen, duck, goose...).
        self.fc2 = nn.Linear(128, num_classes)
    # forward defines the actual PATH the data takes, in order, from when
    # the image enters until the prediction comes out.
    def forward(self, x):
        # Step 1: convolution -> ReLU activation -> pooling
        # ReLU simply turns off (sets to 0) negative values; helps the network learn better.
        x = self.pool(torch.relu(self.conv1(x)))
        # Step 2: we repeat the same pattern with the second convolution
        x = self.pool(torch.relu(self.conv2(x)))
        # Step 3: we "flatten" the tensor. It goes from having shape [channels, height, width]
        # to being a single long vector, which is what a dense (Linear) layer expects.
        # The -1 tells PyTorch "calculate the size of this axis automatically".
        x = x.view(-1, 32 * 8 * 8)
        # Step 4: we go through the first dense layer + ReLU
        x = torch.relu(self.fc1(x))
        # Step 5: output layer. NO activation here on purpose
        # (CrossEntropyLoss, below, already applies softmax internally).
        x = self.fc2(x)
        return x
# Creamos una instancia de la red y la enviamos al dispositivo elegido antes (GPU o CPU).
net = Net(num_classes).to(device)

Every network in PyTorch is defined in two different places, and it is important not to confuse them:

• `_init_` — here you only declare the layers you will use, like someone preparing the tools on the table before starting to work. There is no data flow yet.
• `forward` — here you define the actual order in which those layers process the image. It is the path the data takes from entering until it comes out converted into a prediction.

Layer Functions

A couple of analogies to make these pieces truly clear:

• `padding=1` is like putting a one-pixel frame around the photo. If we didn't do it, the 3×3 filter wouldn't be able to "center" on the outer edge pixels (it would lack a neighbor on one side), and the resulting image would come out slightly smaller each time it passes through a convolution. The added frame avoids this cropping.
• `MaxPool2d(2, 2)` is like reducing a photo to half size by keeping only the most "intense" pixel of each 2×2 block. Imagine you divide the image into squares of 4 pixels and, from each square, only keep the one with the highest value. You lose fine detail, but preserve what "stood out" most in that area — that's why the size is halved on each pass.
• `x.view(-1, 32 * 8 * 8)` is like unrolling a folded blanket into a single long thread. Before this line, the data has the shape of a "box" (32 layers of 8×8 pixels each). A dense (`Linear`) layer does not understand boxes, only lists of numbers in a row. `view()` does not change any values, it just "unrolls" everything into a single long list of 2048 numbers so the dense layer can process it.

Why `32 * 8 * 8`: the image enters at 32×32. Each `MaxPool2d(2,2)` halves the size, and we apply pooling twice (once after each convolution), so 32 → 16 → 8. With 32 filters in the last convolution, the flattened tensor reaching the dense layer has 32 × 8 × 8 = 2048 values. This calculation must be redone every time you change the input size or the number of convolutional layers — it is the most common error when modifying an existing architecture.

4. The Training and Evaluation Loop

# The "loss function" measures how much the model errs in each prediction.
# CrossEntropyLoss is the standard for multi-class classification.
criterion = nn.CrossEntropyLoss()
# The optimizer is what adjusts the model weights to reduce that loss.
# Adam is the most common default choice; lr=0.001 is the "step size" of each adjustment.
optimizer = optim.Adam(net.parameters(), lr=0.001)
# An "epoch" = a complete pass through the ENTIRE training dataset.
# We train for 10 epochs.
for epoch in range(10):
    running_loss = 0.0   # accumulator to see how the error evolves within the epoch
    # train_loader serves us batches of 32 images with their labels.
    for i, (inputs, labels) in enumerate(train_loader):
        # We move the current batch to the same device where the model is (GPU or CPU).
        inputs, labels = inputs.to(device), labels.to(device)
        # PyTorch accumulates gradients by default; they must be reset in each batch
        # or they would sum with those of the previous batch and the training would break.
        optimizer.zero_grad()
        # Forward pass: we pass the images to the model and get its predictions.
        outputs = net(inputs)
        # We compare the prediction with the real label to calculate the error.
        loss = criterion(outputs, labels)
        # Backpropagation: calculates how much each model weight contributed to the error.
        loss.backward()
        # Applies those calculations: slightly adjusts the model weights to reduce the error.
        optimizer.step()
        running_loss += loss.item()   # we save the error value to be able to monitor it

The last layer (`fc2`) does not have its own activation because `CrossEntropyLoss` already applies softmax internally — adding it twice would be a common error that distorts training.

At the end, the model is evaluated against the test set (images it never saw during training) and saved:

# We build the final path where the trained model will be saved.
model_path = os.path.join(MODEL_DIR, "modelo_local.pth")
# state_dict() contains only the learned WEIGHTS (not the architecture).
# Therefore, to use this model again later, you will also need the Net class
# defined exactly as here, and then load these weights into it.
torch.save(net.state_dict(), model_path)

5. What We Have Built So Far

• A project with a standard structure (data / model / script).
• A minimal CNN with two convolutional layers and two dense layers.
• A complete training loop, with its evaluation and weights saving.

What we have not explained yet is what is actually happening inside `conv1` and `conv2` when we say they "detect patterns" — nor what exactly `pool` does beyond "reducing size". That is what we develop in the second installment of this series, where we go into the concept of convolution with visual examples on real images.

6. Full Code

'''
EXAMPLE CODE FOR TRAINING A VISION MODEL:
testENTRENAR_IA/
├── data/
│   └── cifar-10/     # images and labels
├── models/
│   └── modelo_cifar10.pth       # trained weights
└── train_cifar10.py             # training script
'''
# scripts/train_local_images.py
import os
import torch
import torchvision
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
# === LOCAL PATHS ===
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))                  # testENTRENAR_IA/
DATA_DIR = os.path.join(ROOT_DIR, "data", "cifar10")                   # ./data/
TRAIN_DIR = os.path.join(DATA_DIR, "train")
TEST_DIR = os.path.join(DATA_DIR, "test")
MODEL_DIR = os.path.join(ROOT_DIR, "models")
os.makedirs(MODEL_DIR, exist_ok=True)
# === DEVICE ===
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"✅ Dispositivo: {device}")
# === TRANSFORMACIONES ===
transform = transforms.Compose([
transforms.Resize((32, 32)),           # igual que CIFAR-10
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
#transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  #Normalización correcta para RGB
])
# === LOAD DATA ===
train_dataset = ImageFolder(root=TRAIN_DIR, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = ImageFolder(root=TEST_DIR, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
# === NUMBER OF CLASSES ===
num_classes = len(train_dataset.classes)
print(f"📚 Clases detectadas: {train_dataset.classes}")
# === SIMPLE MODEL ===
class Net(nn.Module):
def __init__(self, num_classes):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.fc1 = nn.Linear(32 * 8 * 8, 128)
self.fc2 = nn.Linear(128, num_classes)
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = x.view(-1, 32 * 8 * 8)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
net = Net(num_classes).to(device)
# === TRAINING ===
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
for epoch in range(10):
running_loss = 0.0
for i, (inputs, labels) in enumerate(train_loader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
print(f"[{epoch + 1}, {i + 1}] pérdida: {running_loss / 10:.3f}")
running_loss = 0.0
print("✅ Entrenamiento finalizado")
# === EVALUATION ===
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = net(inputs)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f"📊 Precisión en test: {100 * correct / total:.2f}%")
# === SAVE MODEL ===
model_path = os.path.join(MODEL_DIR, "modelo_local.pth")
torch.save(net.state_dict(), model_path)
print(f"💾 Modelo guardado en: {model_path}")