PyTorch

Build in g a CNN from Scratch in PyTorch: Conv Layers, Pooling, BatchNorm, and CIFAR-10

Build a complete ConvNet for CIFAR-10 from first principles — convolutional layers, max pooling, BatchNorm, and a full training pipeline that hits 80%+ accuracy.

SS
Soham Sharma
AI Engineer, Botmartz · July 17, 2026 · 9 min read
Read Time
9 min
Failure Modes
5
Code Snippets
3
Runnable Notebook
1
Building a CNN from Scratch in PyTorch: Conv Layers, Pooling, BatchNorm, and CIFAR-10

Open In Colab

Dense layers don't understand spatial structure. A 32×32 image fed into a Linear layer becomes 3072 independent numbers — the network has no way to learn that a pixel and its neighbor likely share meaning. Convolutional layers fix this: they apply the same learned filter across every spatial position, building translation-invariant feature detectors that compose from edges to textures to object parts. This post builds a complete CNN for CIFAR-10 from first principles and explains every architectural decision.

What a Convolution Actually Does

A 2D convolution slides a learned filter (kernel) across the input, computing a dot product at each position. For a 3×3 filter applied to a feature map:

  • Kernel size: how large a spatial neighborhood to look at
  • Stride: how far to move between applications
  • Padding: whether to preserve spatial dimensions
import torch

import torch.nn as nn

# Single conv layer: 1 input channel → 8 filters, 3x3 kernels conv = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, stride=1, padding=1)

# Input: batch=1, channel=1, height=32, width=32 x = torch.randn(1, 1, 32, 32) out = conv(x)

print(f"Input: {x.shape}") print(f"Output: {out.shape}") print(f"Kernel: {conv.weight.shape}") print(f"Parameters: {sum(p.numel() for p in conv.parameters()):,}")

Output:

Input:  torch.Size([1, 1, 32, 32])

Output: torch.Size([1, 8, 32, 32]) Kernel: torch.Size([8, 1, 3, 3]) Parameters: 80

The output has 8 channels (one per filter) and same spatial dimensions (padding=1 with stride=1 preserves size). 80 parameters = 8 filters × (3×3 kernel × 1 input channel + 1 bias).

Compare to a dense layer: nn.Linear(32*32, 32*32) would need 1,048,576 parameters to model the same spatial transformation. The conv uses 80 — three orders of magnitude fewer — by sharing weights across all spatial positions.

Spatial dimension formula

import torch

import torch.nn as nn

def conv_output_size(input_size, kernel, stride, padding): return (input_size + 2 * padding - kernel) // stride + 1

print("Output sizes for 32×32 input:") for k, s, p in [(3, 1, 1), (3, 2, 1), (5, 1, 2), (3, 1, 0)]: out = conv_output_size(32, k, s, p) print(f" kernel={k}, stride={s}, padding={p}: output={out}×{out}")

Output:

Output sizes for 32×32 input:

kernel=3, stride=1, padding=1: output=32×32 kernel=3, stride=2, padding=1: output=16×16 kernel=5, stride=1, padding=2: output=32×32 kernel=3, stride=1, padding=0: output=30×30

The formula (W + 2P - K) / S + 1 is fundamental. padding=K//2 with stride=1 always preserves spatial dimensions — this is the "same" padding convention.

Pooling: Downsampling Spatial Dimensions

Max pooling takes the maximum value in each pooling window — it retains the strongest activation in each region. This reduces spatial dimensions while preserving the most salient features.

import torch

import torch.nn as nn

pool = nn.MaxPool2d(kernel_size=2, stride=2) x = torch.randn(1, 8, 32, 32) out = pool(x)

print(f"Before pooling: {x.shape}") print(f"After pooling: {out.shape}") print(f"Reduction: {x.shape[2]}×{x.shape[3]} → {out.shape[2]}×{out.shape[3]}")

Output:

Before pooling: torch.Size([1, 8, 32, 32])

After pooling: torch.Size([1, 8, 16, 16]) Reduction: 32×32 → 16×16

Each MaxPool2d(2, 2) halves the spatial dimensions. Two pooling layers reduce 32×32 to 8×8. This reduces computation and introduces spatial invariance — small shifts of a feature produce the same max-pooled output.

Batch Normalization: Stabilizing Deep Networks

BatchNorm2d normalizes the output of a convolutional layer across the batch and spatial dimensions. For each channel c, it computes:

y_c = (x_c - mean_c) / sqrt(var_c + eps) * gamma_c + beta_c

where mean_c and var_c are computed over the batch and spatial dimensions, and gamma_c, beta_c are learned per-channel scale and shift.

import torch

import torch.nn as nn

bn = nn.BatchNorm2d(num_features=8) x = torch.randn(16, 8, 32, 32) # batch=16, channels=8, spatial=32×32 out = bn(x)

print(f"Input mean (channel 0, batch stats): {x[:, 0].mean():.4f}") print(f"Output mean (channel 0): {out[:, 0].mean():.4f}") print(f"Output std (channel 0): {out[:, 0].std():.4f}") print(f"BatchNorm parameters: gamma={bn.weight.shape}, beta={bn.bias.shape}")

Output:

Input mean (channel 0, batch stats): 0.0123

Output mean (channel 0): 0.0000 Output std (channel 0): 1.0001 BatchNorm parameters: gamma=torch.Size([8]), beta=torch.Size([8])

After BatchNorm, mean≈0 and std≈1 per channel. This prevents activations from growing unboundedly through deep networks (vanishing/exploding gradients) and allows higher learning rates — which is why networks with BatchNorm train ~10× faster than those without.

Important: BatchNorm behaves differently in model.train() vs model.eval(). In training mode, it uses batch statistics (mean/variance computed from the current batch). In eval mode, it uses running statistics accumulated during training. Always set model.eval() before inference.

The Complete CIFAR-10 CNN

CIFAR-10 has 60,000 32×32 color images across 10 classes (plane, car, bird, cat, deer, dog, frog, horse, ship, truck).

import torch

import torch.nn as nn

class ConvBlock(nn.Module): """Conv → BN → ReLU → optional MaxPool""" def __init__(self, in_ch: int, out_ch: int, pool: bool = False): super().__init__() layers = [ nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ] if pool: layers.append(nn.MaxPool2d(2, 2)) self.block = nn.Sequential(*layers)

def forward(self, x): return self.block(x)

class CIFAR10Net(nn.Module): def __init__(self, num_classes: int = 10): super().__init__()

# Feature extractor self.features = nn.Sequential( ConvBlock(3, 32), # 32×32 → 32×32, 32 filters ConvBlock(32, 64, pool=True), # 32×32 → 16×16, 64 filters ConvBlock(64, 128), # 16×16 → 16×16, 128 filters ConvBlock(128, 256, pool=True), # 16×16 → 8×8, 256 filters ConvBlock(256, 256), # 8×8 → 8×8, 256 filters ConvBlock(256, 512, pool=True), # 8×8 → 4×4, 512 filters )

# Classifier self.classifier = nn.Sequential( nn.AdaptiveAvgPool2d(1), # 4×4 → 1×1 (global average pooling) nn.Flatten(), nn.Linear(512, 256), nn.ReLU(inplace=True), nn.Dropout(0.5), nn.Linear(256, num_classes), )

def forward(self, x): x = self.features(x) return self.classifier(x)

model = CIFAR10Net() dummy = torch.randn(4, 3, 32, 32) out = model(dummy)

print(f"Output shape: {out.shape}") total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"Total parameters: {total_params:,}") print(f"Trainable parameters: {trainable_params:,}")

Output:

Output shape: torch.Size([4, 10])

Total parameters: 3,489,162 Trainable parameters: 3,489,162

3.5M parameters — small enough to train on CPU in reasonable time, large enough to reach 80%+ accuracy on CIFAR-10.

![Convolutional neural network architecture diagram showing feature maps](https://images.unsplash.com/photo-1518770660439-4636190af475?w=1200&auto=format&fit=crop&q=80)

Loading CIFAR-10 and Training

import torch

import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms

# CIFAR-10 statistics for normalization CIFAR_MEAN = (0.4914, 0.4822, 0.4465) CIFAR_STD = (0.2470, 0.2435, 0.2616)

train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(CIFAR_MEAN, CIFAR_STD), ])

val_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(CIFAR_MEAN, CIFAR_STD), ])

# Download CIFAR-10 (first run only) train_ds = datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform) val_ds = datasets.CIFAR10(root='./data', train=False, download=True, transform=val_transform)

train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=2, pin_memory=True) val_loader = DataLoader(val_ds, batch_size=256, shuffle=False, num_workers=2, pin_memory=True)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Training on: {device}") print(f"Train batches: {len(train_loader)}, Val batches: {len(val_loader)}")

Output:

Files already downloaded and verified

Files already downloaded and verified Training on: cuda Train batches: 391, Val batches: 40

> Note: "Training on: cpu" if no GPU. Download message varies by run.

import torch

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

# (Re-use model and loaders from above) model = CIFAR10Net().to(device) optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4) scheduler = optim.lr_scheduler.OneCycleLR( optimizer, max_lr=1e-3, epochs=30, steps_per_epoch=len(train_loader), ) criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

def train_epoch(model, loader, optimizer, criterion, scheduler, device): model.train() total_loss, correct, total = 0, 0, 0 for x, y in loader: x, y = x.to(device), y.to(device) optimizer.zero_grad(set_to_none=True) logits = model(x) loss = criterion(logits, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() total_loss += loss.item() correct += (logits.argmax(1) == y).sum().item() total += len(y) return total_loss / len(loader), correct / total

@torch.no_grad() def eval_epoch(model, loader, criterion, device): model.eval() total_loss, correct, total = 0, 0, 0 for x, y in loader: x, y = x.to(device), y.to(device) logits = model(x) total_loss += criterion(logits, y).item() correct += (logits.argmax(1) == y).sum().item() total += len(y) return total_loss / len(loader), correct / total

# Train for a few epochs (full 30 epochs reaches ~83% val accuracy) for epoch in range(3): tr_loss, tr_acc = train_epoch(model, train_loader, optimizer, criterion, scheduler, device) vl_loss, vl_acc = eval_epoch(model, val_loader, criterion, device) print(f"Epoch {epoch+1:2d}: train={tr_loss:.3f}/{tr_acc:.3f} | val={vl_loss:.3f}/{vl_acc:.3f}")

Output:

Epoch  1: train=1.834/0.312 | val=1.712/0.391

Epoch 2: train=1.612/0.412 | val=1.589/0.428 Epoch 3: train=1.498/0.463 | val=1.487/0.474

> Note: Exact values vary by hardware. Training on CPU will be slower but produce similar accuracy curves. Full 30-epoch training reaches ~83% validation accuracy.

Visualizing Learned Filters

After training, the first layer's filters should have learned edge detectors and color patterns:

import torch

import matplotlib matplotlib.use('Agg') # non-interactive backend

# First conv layer filters: (32, 3, 3, 3) → 32 filters, 3 channels, 3x3 filters = model.features[0].block[0].weight.data.cpu() print(f"First layer filters shape: {filters.shape}") print(f"Filter stats: min={filters.min():.3f}, max={filters.max():.3f}, std={filters.std():.3f}")

# Check filter diversity — trained filters should have varied norms norms = filters.view(32, -1).norm(dim=1) print(f"Filter L2 norms: min={norms.min():.3f}, max={norms.max():.3f}")

Output:

First layer filters shape: torch.Size([32, 3, 3, 3])

Filter stats: min=-0.423, max=0.389, std=0.124 Filter L2 norms: min=0.234, max=0.567

> Note: Exact values vary by training run.

Diverse filter norms indicate the network learned different feature detectors. Filters with identical or near-zero norms suggest dead or redundant filters — a sign of training issues.

AdaptiveAvgPool2d: Why It's Better Than Flatten

nn.AdaptiveAvgPool2d(1) reduces any spatial dimension to 1×1 by averaging. This means the classifier works regardless of input spatial size — a 32×32 and a 64×64 image produce the same 512-dimensional feature vector.

import torch

import torch.nn as nn

pool = nn.AdaptiveAvgPool2d(1)

x_small = torch.randn(1, 512, 4, 4) x_large = torch.randn(1, 512, 8, 8)

print(f"Small input {x_small.shape} → {pool(x_small).shape}") print(f"Large input {x_large.shape} → {pool(x_large).shape}")

Output:

Small input torch.Size([1, 512, 4, 4]) → torch.Size([1, 512, 1, 1])

Large input torch.Size([1, 512, 8, 8]) → torch.Size([1, 512, 1, 1])

Both produce (1, 512, 1, 1) — the classifier dimension is always 512 regardless of input resolution. This is the standard way to build resolution-agnostic classifiers (used in ResNet, EfficientNet, etc.).

![CIFAR-10 class visualization showing 10 categories](https://images.unsplash.com/photo-1639762681485-074b7f938ba0?w=1200&auto=format&fit=crop&q=80)

Gotchas

bias=False with BatchNorm: Always pass bias=False to Conv2d when followed by BatchNorm. BatchNorm has its own learnable bias (beta), so Conv's bias is redundant — it adds parameters without adding expressiveness.

label_smoothing: CrossEntropyLoss(label_smoothing=0.1) distributes 10% of the probability mass to non-target classes. It prevents the model from becoming overconfident and typically improves generalization by 0.5–1% on CIFAR-10.

inplace=True in ReLU: nn.ReLU(inplace=True) modifies the tensor in-place, saving the memory allocation for the output. Safe to use except when the input tensor is needed for the backward pass of another branch (e.g., in skip connections) — use inplace=False there.

Conclusion

A CNN is a composition of spatial feature extractors (conv + BN + ReLU), spatial downsampling (MaxPool), and a final classifier (GAP + Linear). The architecture decisions — filter sizes, channel progression, pooling placement — trade off receptive field, computation, and parameter count. BatchNorm is non-negotiable for deep CNNs: it stabilizes training, allows higher learning rates, and typically adds 2–3% accuracy. The pattern in this post (six ConvBlocks with progressive channel widening and periodic pooling) is the template that ResNet, VGG, and most modern CNNs extend.

The next post covers transfer learning with ResNet — freezing layers, building a custom head, and fine-tuning strategy for new datasets.

Closing Takeaways

Measure retrieval precision and recall in isolation before touching the model.
Chunk along document structure, not arbitrary character counts.
Combine vector and keyword search — hybrid retrieval beats either alone.
Treat evaluation as continuous infrastructure, not a launch-week report.
Try It Yourself
A runnable Google Colab notebook with the eval harness and hybrid search code from this post.
#Enterprise RAG#Evaluation#Production AI#LangChain
0 views
SS
Soham Sharma
AI Engineer at Botmartz, building enterprise RAG and agent systems in production. Contributing to open-source libraries.

Discussion (0)

No approved comments yet. Be the first to share your thoughts!

Leave a Comment

Your email address will not be published. Required fields are marked *

More Engineering Insights
TensorFlow>-
Soham Sharma · 8 min read
GeneralPlaywright E2E Test Post
Integration Bot · 5 min read