Let's actually build a neural network — twice

Here's the very first code you'd run. It loads MNIST, normalizes the pixels, and prints the shapes so you understand exactly what you're holding:

import numpy as np
from sklearn.datasets import fetch_openml

# Download MNIST (70,000 images of 784 pixels each)
mnist = fetch_openml('mnist_784', version=1, as_frame=False)
X = mnist.data.astype(np.float32)      # shape (70000, 784), values 0..255
y = mnist.target.astype(np.int64)      # shape (70000,),   values 0..9

# Normalize pixels to the 0..1 range (Session 7: scaled inputs train better)
X = X / 255.0

# Split: 60k train, 10k test (MNIST's standard split is the last 10k)
X_train, X_test = X[:60000], X[60000:]
y_train, y_test = y[:60000], y[60000:]

print(X_train.shape, y_train.shape)   # (60000, 784) (60000,)
print(X_train.min(), X_train.max())   # 0.0 1.0
print("first label:", y_train[0]) # e.g. 5

Notice the normalization step — dividing by 255. Recall from Session 7 that neural nets train far more reliably when inputs are on a small, consistent scale. Feeding raw 0–255 values would make the early gradients wildly large.

This is a full, runnable program. It assumes X_train, y_train from Topic 1. Read the comments — every line maps to something we've studied.

import numpy as np

np.random.seed(42)

# ---------- helpers ----------
def one_hot(y, num_classes=10):
    out = np.zeros((y.size, num_classes))
    out[np.arange(y.size), y] = 1      # put a 1 at the correct digit
    return out

def relu(z):
    return np.maximum(0, z)        # negatives become 0

def softmax(z):
    # subtract row max for numerical stability, then normalize
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)

# ---------- initialize weights (Session 7: He init for ReLU) ----------
n_in, n_hidden, n_out = 784, 128, 10
W1 = np.random.randn(n_in, n_hidden) * np.sqrt(2.0 / n_in)
b1 = np.zeros(n_hidden)
W2 = np.random.randn(n_hidden, n_out) * np.sqrt(2.0 / n_hidden)
b2 = np.zeros(n_out)

# ---------- hyperparameters ----------
lr        = 0.1      # learning rate (Session 3)
epochs    = 20       # full passes over the training data
batch     = 128      # mini-batch size (Session 3: SGD)
n         = X_train.shape[0]
Y_train   = one_hot(y_train)

# ---------- the training loop ----------
for epoch in range(epochs):
    # shuffle each epoch so batches differ
    perm = np.random.permutation(n)
    Xs, Ys = X_train[perm], Y_train[perm]

    for i in range(0, n, batch):
        xb = Xs[i:i+batch]                  # (B, 784)
        yb = Ys[i:i+batch]                  # (B, 10) one-hot
        B  = xb.shape[0]

        # ----- FORWARD PASS -----
        z1 = xb @ W1 + b1                   # (B, 128)
        a1 = relu(z1)                       # (B, 128)
        z2 = a1 @ W2 + b2                   # (B, 10) logits
        a2 = softmax(z2)                    # (B, 10) probabilities

        # ----- LOSS (cross-entropy, averaged over the batch) -----
        loss = -np.sum(yb * np.log(a2 + 1e-9)) / B

        # ----- BACKPROP (chain rule, Session 5) -----
        dz2 = (a2 - yb) / B                 # softmax+CE shortcut: pred - actual
        dW2 = a1.T @ dz2                    # (128, 10)
        db2 = dz2.sum(axis=0)            # (10,)

        da1 = dz2 @ W2.T                    # (B, 128)
        dz1 = da1 * (z1 > 0)              # ReLU gradient: 1 where z1>0 else 0
        dW1 = xb.T @ dz1                    # (784, 128)
        db1 = dz1.sum(axis=0)            # (128,)

        # ----- GRADIENT-DESCENT UPDATE (Session 3) -----
        W2 -= lr * dW2;  b2 -= lr * db2
        W1 -= lr * dW1;  b1 -= lr * db1

    print(f"epoch {epoch+1:2d}  loss {loss:.4f}")

# ---------- evaluate on the held-out test set ----------
def predict(X):
    a1 = relu(X @ W1 + b1)
    return softmax(a1 @ W2 + b2).argmax(axis=1)

acc = (predict(X_test) == y_test).mean()
print(f"test accuracy: {acc*100:.2f}%")   # ~97% after 20 epochs

That is a complete neural network — forward, loss, backprop, update — in about 50 lines, with no framework. Run it and you'll watch the loss fall and finish around 97% test accuracy.

A full, runnable training script. Compare it line-by-line with the NumPy version — the structure is identical, but backprop is now a single call.

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

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

# ---------- 1. DATA ----------
tfm = transforms.Compose([
    transforms.ToTensor(),                      # to tensor, scales pixels to 0..1
    transforms.Normalize((0.1307,), (0.3081,)),  # MNIST mean/std
])
train_ds = datasets.MNIST("./data", train=True,  download=True, transform=tfm)
test_ds  = datasets.MNIST("./data", train=False, download=True, transform=tfm)

train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
test_loader  = DataLoader(test_ds,  batch_size=256)

# ---------- 2. MODEL ----------
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),               # (B,1,28,28) -> (B,784)
            nn.Linear(784, 128),       # hidden layer
            nn.ReLU(),                  # activation (Session 4)
            nn.Linear(128, 10),        # output layer (logits)
        )
    def forward(self, x):
        return self.net(x)

model = MLP().to(device)

# ---------- 3. LOSS + OPTIMIZER ----------
# CrossEntropyLoss applies softmax internally, so the model outputs raw logits
loss_fn   = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# ---------- 4. TRAIN ----------
for epoch in range(20):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)

        logits = model(xb)              # FORWARD
        loss   = loss_fn(logits, yb)    # LOSS

        optimizer.zero_grad()           # clear last step's gradients (!)
        loss.backward()                 # BACKPROP — autograd does it all
        optimizer.step()                # UPDATE every weight
    print(f"epoch {epoch+1:2d}  loss {loss.item():.4f}")

# ---------- 5. EVALUATE ----------
model.eval()
correct = 0
with torch.no_grad():                   # no gradients needed at test time
    for xb, yb in test_loader:
        xb, yb = xb.to(device), yb.to(device)
        preds = model(xb).argmax(dim=1)
        correct += (preds == yb).sum().item()
print(f"test accuracy: {100*correct/len(test_ds):.2f}%")  # ~97-98%

Same architecture, same result — but the hand-written gradients are gone. The three lines zero_grad() → backward() → step() are the heart of every PyTorch training loop you will ever write.

Here's the practical loop you'd actually use, computing validation loss each epoch:

def evaluate(loader):
    model.eval()
    total, correct, loss_sum = 0, 0, 0.0
    with torch.no_grad():
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            logits = model(xb)
            loss_sum += loss_fn(logits, yb).item() * xb.size(0)
            correct  += (logits.argmax(1) == yb).sum().item()
            total    += xb.size(0)
    return loss_sum / total, correct / total

best_val, patience, waited = 1e9, 3, 0
for epoch in range(50):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        loss = loss_fn(model(xb), yb)
        optimizer.zero_grad(); loss.backward(); optimizer.step()

    val_loss, val_acc = evaluate(val_loader)
    print(f"epoch {epoch+1}  val_loss {val_loss:.4f}  val_acc {val_acc*100:.2f}%")

    # EARLY STOPPING: stop if val loss hasn't improved for `patience` epochs
    if val_loss < best_val:
        best_val, waited = val_loss, 0
        torch.save(model.state_dict(), "best.pt")  # keep the best model
    else:
        waited += 1
        if waited >= patience:
            print("early stopping"); break

The early stopping logic — save the model whenever validation loss hits a new low, and quit if it hasn't improved in a few epochs — is the simplest, most effective regularizer there is.