Python
# train.py — trains and saves an MNIST CNN classifier
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

device = "cuda" if torch.cuda.is_available() else "cpu"

class MNISTCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.flatten(1)
        x = self.dropout(F.relu(self.fc1(x)))
        return self.fc2(x)

def main():
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    train_ds = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
    test_ds = datasets.MNIST(root="./data", train=False, download=True, transform=transform)
    train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
    test_loader = DataLoader(test_ds, batch_size=256)

    model = MNISTCNN().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(5):
        model.train()
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            loss = F.cross_entropy(model(x), y)
            loss.backward()
            optimizer.step()

        model.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for x, y in test_loader:
                x, y = x.to(device), y.to(device)
                pred = model(x).argmax(1)
                correct += (pred == y).sum().item()
                total += y.size(0)
        print(f"Epoch {epoch+1}: test accuracy = {correct/total:.4f}")

    torch.save(model.state_dict(), "mnist_cnn.pt")
    print("Saved model to mnist_cnn.pt")

if __name__ == "__main__":
    main()

Run: python train.py → produces mnist_cnn.pt.

Progress:

• Step 1: Set up environment (torch, torchvision)
• Step 2: Load MNIST via torchvision.datasets.MNIST (auto-downloads)
• Step 3: Define CNN architecture (2 conv + 2 pool + 2 fc layers is sufficient)
• Step 4: Train for 5 epochs with Adam optimizer, cross-entropy loss
• Step 5: Evaluate on test split each epoch, print accuracy
• Step 6: Save model weights (torch.save(model.state_dict(), ...))
• Step 7: Write separate infer.py for inference on new images
• Step 8: Validate inference on a few test samples before trusting on custom images

Inference script

Python
# infer.py — loads trained model and predicts on an image
import sys
import torch
from torchvision import transforms
from PIL import Image
from train import MNISTCNN  # reuse architecture definition

device = "cuda" if torch.cuda.is_available() else "cpu"

def load_model(path="mnist_cnn.pt"):
    model = MNISTCNN().to(device)
    model.load_state_dict(torch.load(path, map_location=device))
    model.eval()
    return model

def preprocess(image_path):
    transform = transforms.Compose([
        transforms.Grayscale(num_output_channels=1),
        transforms.Resize((28, 28)),
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    img = Image.open(image_path)
    return transform(img).unsqueeze(0).to(device)  # add batch dim

def predict(image_path, model_path="mnist_cnn.pt"):
    model = load_model(model_path)
    x = preprocess(image_path)
    with torch.no_grad():
        logits = model(x)
        probs = torch.softmax(logits, dim=1)
        pred = probs.argmax(1).item()
    return pred, probs.squeeze().tolist()

if __name__ == "__main__":
    image_path = sys.argv[1]
    digit, confidences = predict(image_path)
    print(f"Predicted digit: {digit}")
    print(f"Confidence: {confidences[digit]:.4f}")

Run: python infer.py path/to/digit.png