DAI developerLab· about 60 min· server sandbox

Lab: train a digit classifier in PyTorch

Build an nn.Module, a DataLoader, a training loop with CrossEntropyLoss and evaluation on held-out data — and reach ≥ 92 % on 8×8 digits.

Teaches: Train a neural network in PyTorch

Requires: PyTorch — tensors and autogradBackpropagation

Theory

Zero, forward, loss, backward, step. CrossEntropyLoss takes logits. Validate on data the model has not seen.

Sub-tasks

the model — make_model() returns an nn.Sequential 64 → 64 → 10 with ReLU.
one training step — train_step(model, opt, loss_fn, X, y) takes one step and returns the loss as a float.
train and evaluate — fit(epochs, lr, seed) trains on sklearn digits (80 % train) and returns the validation accuracy.

Passes when: accuracy >= 0.92

The starter code

runs in an isolated sandbox on the server

import torch
import torch.nn as nn
from sklearn.datasets import load_digits


def data(seed=0):
    d = load_digits()
    X = torch.tensor(d.data, dtype=torch.float32) / 16.0
    y = torch.tensor(d.target)
    g = torch.Generator().manual_seed(seed)
    idx = torch.randperm(len(X), generator=g)
    n = int(0.8 * len(X))
    return X[idx[:n]], y[idx[:n]], X[idx[n:]], y[idx[n:]]


def make_model():
    # TODO: nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 10))
    ...


def train_step(model, opt, loss_fn, X, y):
    # TODO: zero_grad → forward → loss → backward → step; return loss.item()
    ...


def accuracy(model, X, y):
    model.eval()
    with torch.no_grad():
        return float((model(X).argmax(1) == y).float().mean())


def fit(epochs=20, lr=1e-2, seed=0, batch=64):
    torch.manual_seed(seed)
    Xtr, ytr, Xva, yva = data(seed)
    model = make_model()
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    # TODO: loop över epoker och batchar (torch.randperm för blandning), anropa train_step
    ...
    return accuracy(model, Xva, yva)

You write the code; tests you cannot see decide whether it holds up. Create a free account to run the lab.

Try the diagnostic Create a free account

Expected results

Validation accuracy ≥ 0.92 after ~20 epochs with Adam lr 1e-2 (often 0.96+).

Common mistakes

Softmax before CrossEntropyLoss (double softmax) → slow training.
Evaluates on training data → falsely high accuracy.
Forgets model.eval()/torch.no_grad() during evaluation.