Building and Training CNNs in PyTorch
nn.Conv2d, nn.MaxPool2d, nn.Linear, training loop, optimizer, loss, accuracy tracking.
Building and Training CNNs in PyTorch is a free Learn AI with Python lesson on CoddyKit — lesson 3 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the Learn AI with Python learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.
Why CNNs for Images?
Convolutional Neural Networks learn spatial patterns: edges, textures, then shapes and objects. Convolutions share weights across the image, making CNNs efficient and translation-aware, the backbone of computer vision.
import torch
import torch.nn as nnThe nn.Module Base Class
Models subclass nn.Module. You define layers in __init__ and the data flow in forward. PyTorch tracks parameters and gradients automatically.
class CNN(nn.Module):
def __init__(self):
super().__init__()Convolution Layers
nn.Conv2d(in_channels, out_channels, kernel_size) slides learnable filters over the image to produce feature maps. The first conv takes 3 channels (RGB) and outputs more channels capturing different features.
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)Pooling Layers
nn.MaxPool2d downsamples feature maps by keeping the max in each window. This shrinks spatial size, reduces computation, and adds a little translation invariance.
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)Fully Connected Layers
After convolutions, flatten the features and pass them through nn.Linear layers to produce class scores. The final Linear outputs one value per class.
self.fc1 = nn.Linear(32 * 8 * 8, 128)
self.fc2 = nn.Linear(128, 10) # 10 classesThe forward Method
The forward method defines how data flows: conv to ReLU to pool, repeated, then flatten and feed the linear layers. F.relu adds the nonlinearity that lets the network learn complex patterns.
import torch.nn.functional as F
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(x.size(0), -1) # flatten
x = F.relu(self.fc1(x))
return self.fc2(x)Loss and Optimizer
For classification use CrossEntropyLoss. An optimizer like Adam updates weights using gradients. Pass the model parameters and a learning rate.
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)The Training Loop: Zero Gradients
Each step starts by clearing old gradients with optimizer.zero_grad(). Skip this and gradients accumulate across batches, corrupting your updates.
for images, labels in loader:
optimizer.zero_grad()Forward Pass and Loss
Run the batch through the model to get predictions, then compute the loss comparing predictions to the true labels. The loss is a single number measuring how wrong the model is.
outputs = model(images)
loss = criterion(outputs, labels)Backward and Step
loss.backward() computes gradients via autograd, and optimizer.step() nudges the weights to reduce the loss. Together they are one learning step.
loss.backward()
optimizer.step()Tracking Accuracy
Monitor progress by counting correct predictions. Take the argmax of the outputs to get predicted classes, compare to labels, and divide hits by the total.
preds = outputs.argmax(dim=1)
correct = (preds == labels).sum().item()
acc = correct / labels.size(0)
print("batch acc:", acc)Quick Check
Test your training loop knowledge.
Recap: Building and Training CNNs
You built a CNN by subclassing nn.Module with nn.Conv2d, nn.MaxPool2d, and nn.Linear layers and a forward method. You trained it with the loop: zero_grad, forward pass, loss.backward(), optimizer.step(), and tracked accuracy via argmax.
Frequently asked questions
Is the “Building and Training CNNs in PyTorch” lesson free?
Yes — the full text of “Building and Training CNNs in PyTorch” is free to read here on the web, and the Learn AI with Python course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the Learn AI with Python course, upgrade to CoddyKit PRO.
What will I learn in “Building and Training CNNs in PyTorch”?
nn.Conv2d, nn.MaxPool2d, nn.Linear, training loop, optimizer, loss, accuracy tracking. You practise Learn AI with Python with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.
Do I need any experience to start Learn AI with Python?
No prior experience is required. Learn AI with Python on CoddyKit is structured for beginners through advanced learners; this is — lesson 3 of 4, so you can start here or from the beginning and move at your own pace.
How long does the “Building and Training CNNs in PyTorch” lesson take?
Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.
Can I write and run code in this Learn AI with Python lesson?
Yes. Every Learn AI with Python lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.
All lessons in this course
- PyTorch Tensors and Autograd
- Custom Datasets and DataLoaders
- Building and Training CNNs in PyTorch
- Object Detection with YOLOv8