What tensor shape does MNIST use in PyTorch?

Batches are [N, 1, 28, 28]. Labels for CrossEntropyLoss are class indices [N], not one-hot.

Why is my first Linear layer the wrong size?

Compute C×H×W after the last pool, then set in_features. A wrong flatten is the usual crash.

Building a Handwritten Digit Recognizer with PyTorch CNNs: From Scratch to 99% Accuracy

The MNIST Handwritten Digit Dataset (70,000 28 \times 28 grayscale images) is the canonical benchmark in computer vision. While a simple Multi-Layer Perceptron can achieve ~96% accuracy, flattening a 2D image into a 1D vector completely discards spatial neighborhood relationships between adjacent pixels.

Convolutional Neural Networks (CNNs) preserve 2D topological spatial structure by sliding learnable convolutional kernels across image feature maps.

In this tutorial, we derive the exact spatial dimension math of 2D convolutions, build a custom CNN in PyTorch, and achieve >99.0% test accuracy.

#1. The Mathematics of 2D Convolution & Pooling

Given an input feature map of height/width W, kernel size K, padding P, and stride S, the output spatial dimension O is governed by:

O = \left\lfloor \frac{W - K + 2P}{S} \right\rfloor + 1

CODE

Input Image [28x28] ──► [Conv2d 3x3, P=1, S=1] ──► Output [28x28]
                          │
                       [MaxPool2d 2x2, S=2]    ──► Output [14x14]
                          │
                       [Conv2d 3x3, P=1, S=1] ──► Output [14x14]
                          │
                       [MaxPool2d 2x2, S=2]    ──► Output [7x7]

#2. Step-by-Step PyTorch CNN Pipeline

Step 1 & 2: TorchVision Ingestion & Normalization

python

import torch
import torchvision
import torchvision.transforms as transforms

# MNIST mean=0.1307, std=0.3081 for global pixel standardization
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_set = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
val_set = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=64, shuffle=False)

Step 3 & 4: CNN Architecture Definition (nn.Module)

python

import torch.nn as nn
import torch.nn.functional as F

class DigitCNN(nn.Module):
    def __init__(self):
        super(DigitCNN, self).__init__()
        # Block 1: 1 channel -> 32 channels (28x28 -> 14x14)
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.pool1 = nn.MaxPool2d(2, 2)
        
        # Block 2: 32 channels -> 64 channels (14x14 -> 7x7)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool2 = nn.MaxPool2d(2, 2)
        
        # Dense Classification Head
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.drop = nn.Dropout(0.3)
        self.fc2 = nn.Linear(128, 10)  # 10 digit classes (0-9)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pool1(F.relu(self.bn1(self.conv1(x))))
        x = self.pool2(F.relu(self.bn2(self.conv2(x))))
        x = x.view(x.size(0), -1)  # Flatten 64*7*7 to 3136
        x = F.relu(self.fc1(x))
        x = self.drop(x)
        x = self.fc2(x)  # Raw unnormalized logits
        return x

model = DigitCNN()

Step 5 & 6: Training with CrossEntropyLoss and Adam

python

import torch.optim as optim

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(1, 6):
    model.train()
    running_loss, correct, total = 0.0, 0, 0
    
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item() * images.size(0)
        _, preds = outputs.max(1)
        total += labels.size(0)
        correct += preds.eq(labels).sum().item()
        
    print(f"Epoch [{epoch}/5] Loss: {running_loss/total:.4f} | Train Acc: {100.0 * correct / total:.2f}%")

Step 7: Test Set Accuracy & JIT TorchScript Export

python

model.eval()
correct, total = 0, 0

with torch.no_grad():
    for images, labels in val_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, preds = outputs.max(1)
        val_total += labels.size(0)
        val_correct += predicted.eq(labels).sum().item()

test_acc = 100.0 * correct / total
print(f"Final MNIST Test Accuracy: {test_acc:.2f}%")

# Export TorchScript for high-performance C++ or browser serving
scripted_model = torch.jit.script(model.to('cpu'))
scripted_model.save("mnist_cnn_jit.pt")
print("TorchScript model exported successfully!")

#Launch the Interactive MNIST Lab

Build and test this exact CNN digit classifier in our browser sandbox with instant AI step verification.

PyTorch CNN MNIST Pipeline from Scratch (7-Step Lab)

PRACTICE THIS PYTORCH PIPELINE IN YOUR BROWSER

Building a Handwritten Digit Recognizer with PyTorch CNNs: From Scratch to 99% Accuracy

#1. The Mathematics of 2D Convolution & Pooling

#2. Step-by-Step PyTorch CNN Pipeline

Step 1 & 2: TorchVision Ingestion & Normalization

Step 3 & 4: CNN Architecture Definition (nn.Module)

Step 5 & 6: Training with CrossEntropyLoss and Adam

Step 7: Test Set Accuracy & JIT TorchScript Export

#Launch the Interactive MNIST Lab

FAQ

What tensor shape does MNIST use in PyTorch?

Why is my first Linear layer the wrong size?

RUN THE PIPELINE IN THE INTERACTIVE LAB