Demo Building and Training a model

Congratulations — you’ve reached the hands-on portion. In this lesson we build and train a simple convolutional neural network (CNN) in PyTorch and walk through standard training and validation loops. Quick review: in PyTorch you define a model by creating a class that inherits from nn.Module. The class typically implements two methods:

__init__ — define layers and components of the network.
forward — define the forward pass through those layers to produce predictions.

Below is a compact, commonly used example: a simple CNN for image classification.

# Imports required for model definition, training, and optimization
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class ImageClassificationNet(nn.Module):
    def __init__(self):
        super(ImageClassificationNet, self).__init__()
        # Input: 1 channel (e.g., grayscale) -> Output: 6 feature maps, kernel 5x5
        self.conv1 = nn.Conv2d(1, 6, 5)
        # Input: 6 feature maps -> Output: 16 feature maps, kernel 5x5
        self.conv2 = nn.Conv2d(6, 16, 5)
        # Max pooling layer to downsample feature maps by factor of 2
        self.pool = nn.MaxPool2d(2, 2)
        # Fully connected layers:
        # After two conv + pool layers on a 32x32 input, the spatial dims are typically reduced,
        # here we assume the flattened conv output has 16 * 5 * 5 = 400 features.
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        # Final layer for 10-class classification
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Conv -> ReLU -> MaxPool
        x = self.pool(F.relu(self.conv1(x)))
        # Conv -> ReLU -> MaxPool
        x = self.pool(F.relu(self.conv2(x)))
        # Flatten while preserving batch dimension
        x = torch.flatten(x, 1)
        # Fully connected layers with ReLU activations
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        # Output raw scores (logits)
        x = self.fc3(x)
        return x

Create a model instance and move it to the appropriate device (CPU or GPU). PyTorch makes device selection straightforward:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = ImageClassificationNet().to(device)
print(model)

Example printed model structure (console output):

ImageClassificationNet(
  (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0)
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)

Tip: Always confirm that your model and data are on the same device. Moving only the model to GPU and forgetting to move the inputs will cause runtime errors.

Loss, optimizer, and hyperparameters

For multi-class classification use CrossEntropyLoss. A simple optimizer to start with is SGD with momentum.

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

A small example hyperparameter set for this demo:

Hyperparameter	Example value	Notes
Learning rate	`0.001`	Reduce if loss diverges; increase slightly if training is too slow.
Momentum	`0.9`	Accelerates SGD in the relevant direction.
Epochs	`3`	Demo uses few epochs — increase for real training.
Batch size	`32` (typical)	Larger batches often require larger learning rates.

N_EPOCHS = 3

Data loaders

Make sure you have train_loader and val_loader prepared as torch.utils.data.DataLoader instances. Each loader should yield batches in the form (inputs, labels). The training and validation loops below assume those are available.

Minimal training loop (no validation)

Typical training loop steps:

model.train() to enable training-time behaviors (dropout, batchnorm updates).
Iterate batches: zero gradients, forward pass, compute loss, backward pass, optimizer step.
Accumulate running loss to monitor progress.

for epoch in range(N_EPOCHS):
    running_loss = 0.0
    model.train()  # set model to training mode

    for i, data in enumerate(train_loader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    avg_loss = running_loss / len(train_loader)
    print(f"Epoch: {epoch} Loss: {avg_loss}")

Sample output (illustrative):

Epoch: 0 Loss: 1.4917420816898346
Epoch: 1 Loss: 0.5658748598178228
Epoch: 2 Loss: 0.48233314801454547

Training with validation

Validation helps monitor generalization and spot overfitting. For validation:

Use model.eval() to disable dropout and use running statistics for batchnorm.
Wrap the validation loop in torch.no_grad() to avoid computing gradients.
Do not call loss.backward() or optimizer.step() during validation.

Combined training + validation loop:

for epoch in range(N_EPOCHS):
    ###### TRAINING ######
    training_loss = 0.0
    model.train()
    for i, data in enumerate(train_loader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        training_loss += loss.item()

    ###### VALIDATION ######
    val_loss = 0.0
    model.eval()
    with torch.no_grad():
        for i, data in enumerate(val_loader, 0):
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)

            outputs = model(inputs)
            loss = criterion(outputs, labels)
            val_loss += loss.item()

    print(
        f"Epoch: {epoch} Train Loss: {training_loss/len(train_loader)} "
        f"Val Loss: {val_loss/len(val_loader)}"
    )

Example output (illustrative):

Epoch: 0 Train Loss: 0.43306035616199173 Val Loss: 0.423773640546555
Epoch: 1 Train Loss: 0.4004011395931244 Val Loss: 0.40547372574528184
Epoch: 2 Train Loss: 0.37513792911569277 Val Loss: 0.38084139679662715

Validation loss decreasing along with training loss is a good sign of generalization. If training loss keeps decreasing but validation loss plateaus or increases, that indicates overfitting — consider regularization, data augmentation, or early stopping.

Notes and best practices

Always move inputs and labels to the same device as the model: inputs = inputs.to(device); labels = labels.to(device).
Use model.train() when training to enable dropout and batchnorm updates; use model.eval() when validating or testing to disable them.
Wrap validation/testing code in with torch.no_grad(): to save memory and computation.
The demo uses N_EPOCHS = 3 for brevity — production models typically require many more epochs.
To improve performance consider:
- Increasing epochs and using early stopping based on validation loss.
- Adjusting learning rate and using learning rate schedules (e.g., ReduceLROnPlateau, StepLR).
- Using data augmentation and regularization (dropout, weight decay).
- Switching to adaptive optimizers (Adam) and monitoring their behavior compared to SGD with momentum.
- Monitoring additional metrics such as accuracy, precision, recall, or confusion matrix for classification.

Quick reference table

Topic	Recommendation
Device	Use GPU if available: `torch.cuda.is_available()`
Loss	`nn.CrossEntropyLoss()` for multi-class classification
Optimizer	Start with `optim.SGD(..., momentum=0.9)` or try `optim.Adam`
Validation	Use `model.eval()` + `torch.no_grad()`
Debugging	Print batch shapes and sample predictions to confirm pipeline correctness

Links and references

You now have a basic, repeatable workflow for defining a PyTorch model, training it, and validating its performance.

Getting Started with PyTorch

Working with Data

Building and Training Models

Model Deployment and Inference

Demo Building and Training a model

Loss, optimizer, and hyperparameters

Data loaders

Minimal training loop (no validation)

Training with validation

Notes and best practices

Quick reference table

Links and references

Watch Video

Practice Lab

​Loss, optimizer, and hyperparameters

​Data loaders

​Minimal training loop (no validation)

​Training with validation

​Notes and best practices

​Quick reference table

​Links and references

Watch Video

Practice Lab

Loss, optimizer, and hyperparameters

Data loaders

Minimal training loop (no validation)

Training with validation

Notes and best practices

Quick reference table

Links and references