Table of Contents
From Idea to `nn.Module`
A basic convolutional neural network in PyTorch is usually built from a small stack of convolution blocks that turn an image tensor into a set of learned feature maps, followed by a classifier head that converts those features into class scores. In code, you package this into an nn.Module so it can be moved across devices, participate in autograd, and expose parameters to an optimizer.
The input you will typically feed to a CNN has shape $(N, C, H, W)$, where $N$ is batch size, $C$ is channels, and $H, W$ are height and width. For example, RGB images are usually $C=3$. Your model’s forward pass should accept that tensor and return logits of shape $(N, \text{num\_classes})$ for classification.
A Minimal CNN Architecture
A common beginner friendly template is two convolution blocks plus a small fully connected head. Each block is usually a convolution plus a nonlinearity, and often some spatial downsampling so the feature maps shrink as depth increases.
Here is a compact, standard implementation:
import torch
import torch.nn as nn
import torch.nn.functional as F
class BasicCNN(nn.Module):
def __init__(self, in_channels: int = 3, num_classes: int = 10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(in_channels, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2), # H,W -> H/2,W/2
nn.Conv2d(16, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2), # H,W -> H/4,W/4
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(32 * 8 * 8, 128), # assumes input images are 32x32
nn.ReLU(),
nn.Linear(128, num_classes)
)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return xThis example assumes your input images are $32 \times 32$. After two $2 \times 2$ max pools, the spatial size becomes $8 \times 8$, so the flattened feature vector has length $32 \cdot 8 \cdot 8$.
:::danger :::
Your first Linear layer must match the flattened size coming out of the conv stack. If you change the number of pools, strides, or the input image size, you must update this dimension or compute it dynamically.
Making the Classifier Head Robust to Input Size
Beginners often run into shape mismatches when they switch datasets with different image sizes. You can avoid hard coding the spatial size by using global pooling that reduces each channel to a single number. In PyTorch, AdaptiveAvgPool2d((1, 1)) converts any $(N, C, H, W)$ into $(N, C, 1, 1)$.
class BasicCNNAdaptive(nn.Module):
def __init__(self, in_channels: int = 3, num_classes: int = 10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(in_channels, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.classifier = nn.Linear(64, num_classes)
def forward(self, x):
x = self.features(x)
x = self.pool(x)
x = torch.flatten(x, 1) # (N, 64)
x = self.classifier(x)
return xThis version works for many input sizes without modification, as long as the feature extractor produces a reasonable spatial resolution before the adaptive pool.
:::danger :::
For multi class classification, return raw logits, not softmax probabilities. Use nn.CrossEntropyLoss, which applies log_softmax internally.
Sanity Checking Shapes with a Dummy Forward Pass
Before training, you should run a single forward pass with a fake batch and confirm that the output shape matches what your loss expects.
model = BasicCNNAdaptive(in_channels=3, num_classes=10)
x = torch.randn(4, 3, 64, 64) # batch of 4 images
logits = model(x)
print(logits.shape) # should be torch.Size([4, 10])If you see a runtime error about matrix multiplication sizes in a linear layer, it means your flatten size and linear input size disagree, or you forgot to flatten.
A Complete Minimal Training Step for the CNN
The training loop structure is covered elsewhere, but it helps to see how the CNN plugs into it. The key is that the CNN outputs logits, you compute loss against integer class labels, then run backprop and an optimizer step.
model = BasicCNNAdaptive(in_channels=3, num_classes=10)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
x = torch.randn(8, 3, 64, 64) # images
y = torch.randint(0, 10, (8,)) # class indices
model.train()
optimizer.zero_grad()
logits = model(x)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
:::danger ::>
CrossEntropyLoss expects targets of shape (N,) with dtype torch.long containing class indices in [0, num_classes-1]. Do not one hot encode the labels for this loss.
Small Practical Implementation Choices
A basic CNN often improves if you add normalization after convolutions. A typical pattern is Conv2d -> BatchNorm2d -> ReLU. If you later add batch normalization, keep in mind that model mode matters, since model.train() and model.eval() change how batch norm behaves.
Another simple improvement is to keep the model definition readable by separating feature extraction and classification into named submodules, as shown above, and keeping forward short so that shape flow is easy to track.
Common Beginner Errors Specific to Basic CNNs
One frequent issue is swapping tensor layout. Many image tools produce (H, W, C) arrays, but Conv2d expects (N, C, H, W). Your preprocessing should ensure channels come before height and width, and you should add a batch dimension.
Another common issue is mismatched channel counts. If your dataset is grayscale, your input has $C=1$, so you must set in_channels=1 in the first Conv2d or convert grayscale to three channels during preprocessing.
:::danger :::
If the first convolution is defined with in_channels=3 but your input batch is (N, 1, H, W), PyTorch will raise a clear error. Fix it by matching in_channels to the data, not by reshaping channels incorrectly.
What You Have After This Chapter
At this point you have a working CNN nn.Module that accepts image batches, produces logits, can be sanity checked with dummy inputs, and can be trained with standard PyTorch losses and optimizers. The next steps are usually improving data and training quality, which are handled in the surrounding chapters on augmentation, transfer learning, and fine tuning.