48.3 Train a Small CNN

By Ravindra Bagale Last updated: 24 October 2026

Why. Reuse the TinyCNN idea and train a few epochs on CPU.

What. CrossEntropyLoss + Adam on the synthetic set.

How.

Train briefly

Define TinyCNN (Chapter 47).
Make train batch of 128.
Run 15–30 epochs; print loss/acc every few steps.

import torch
import torch.nn as nn
import torch.optim as optim

class TinyCNN(nn.Module):
 def __init__(self):
 super().__init__()
 self.features = nn.Sequential(
 nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
 nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
 )
 self.head = nn.Linear(16 * 7 * 7, 2)
 def forward(self, x):
 x = self.features(x)
 return self.head(x.view(x.size(0), -1))

def make_batch(n=128):
 x = torch.rand(n, 1, 28, 28) * 0.2
 y = torch.randint(0, 2, (n,))
 for i in range(n):
 if y[i] == 0: x[i, 0, :, :14] += 0.8
 else: x[i, 0, :, 14:] += 0.8
 return x, y

torch.manual_seed(0)
model = TinyCNN()
opt = optim.Adam(model.parameters(), lr=0.01)
crit = nn.CrossEntropyLoss()
X, y = make_batch()
for epoch in range(20):
 opt.zero_grad()
 logits = model(X)
 loss = crit(logits, y)
 loss.backward()
 opt.step()
 if epoch % 5 == 0:
 acc = (logits.argmax(1) == y).float().mean()
 print(epoch, float(loss), float(acc))

What you should see. Accuracy climbing on this easy synthetic pattern.