Train and validate an MLP without leaking

Two builds share one loop: forward, loss, zero_grad, backward, step on train; forward-only under no_grad on validation. The order below is the build order.

flowchart LR
    ARCH[define arch] --> INIT[init weights]
    INIT --> SPLIT[stratified split]
    SPLIT --> TRAIN[train loop]
    TRAIN --> VALID[validate no-grad]
    VALID --> STOP[stop on best]

Same loop twice: churn teaches the binary setup, FashionMNIST teaches the 10-way setup. Validation never updates weights.

1 · The words this page uses

Seven groups of words carry the rest of the page. Read them once. Each comes back with a number later.

1.1 · Logit, softmax, cross-entropy

A logit is the raw score a Linear layer writes. It can be any real number, so it is not a probability. Softmax turns a row of logits into probabilities that sum to 1: \(p_i = e^{z_i} / \sum_j e^{z_j}\). For two classes one logit is enough: sigmoid, \(\sigma(z) = 1/(1+e^{-z})\), gives the positive-class probability and \(1 - \sigma(z)\) the other. Cross-entropy is the loss that punishes a low probability on the true class: for one row it is \(-\log p_{\text{true}}\). CrossEntropyLoss applies log-softmax (the log of the softmax) and then cross-entropy, so it takes raw logits. BCEWithLogitsLoss does the same pairing for one binary logit: sigmoid plus binary cross-entropy, fused for stability.

1.2 · Epoch, batch, learning rate

An epoch is one pass over the training split. A batch is the group of rows that share one forward pass, one loss, and one weight update. Churn uses the whole 4,508-row training split as one batch per epoch; FashionMNIST uses 64 rows per batch, so 938 batches make one epoch. The learning rate is the size of the step the optimizer takes along the gradient. Too large overshoots the minimum; too small wastes epochs. Adam adapts the step per parameter, but the learning rate still sets its scale. Both models are MLPs (multi-layer perceptrons): stacks of Linear layers, each followed by an activation. ReLU, used here, keeps the positive part of its input and sets negatives to zero.

1.3 · Training, validation, test

The training split fits the weights. The validation split chooses: it triggers early stopping and the learning-rate drop, and it never sends a gradient. The test split is the final check, read once after all choices are made. Churn starts with 7,043 rows: 80% (5,634) for train plus validation and 20% (1,409) for test; the 5,634 are split again into 4,508 train and 1,126 validation. Stratify keeps the churn ratio in every piece, so a rare class cannot vanish from validation.

1.4 · BatchNorm, dropout, train and eval mode

BatchNorm normalizes each feature across the current batch: subtract the batch mean, divide by the batch standard deviation, then apply a learned scale and shift. In train() it uses the batch's own statistics, which change with every batch, and updates a running average. In eval() it uses those running averages, so a row's output does not depend on the other rows in the batch. Dropout has the same switch: in train() it zeroes each hidden unit with probability 0.2, and in eval() it passes every unit. Both layers make model.eval() necessary: the same row must get the same score on every pass.

1.5 · Xavier initialization

Xavier initialization draws each Linear weight from a distribution with variance \(2/(n_{\text{in}} + n_{\text{out}})\), where \(n_{\text{in}}\) is the number of inputs to the layer and \(n_{\text{out}}\) the number of outputs, so the forward signal keeps roughly the same spread from layer to layer. The uniform form samples from \(\pm\sqrt{6/(n_{\text{in}} + n_{\text{out}})}\), whose variance is the same \(2/(n_{\text{in}} + n_{\text{out}})\). Biases start at zero, which removes an arbitrary offset. BatchNorm and Dropout hold no weight matrix, so they are skipped.

1.6 · Patience and early stopping

Patience is the number of epochs a run waits for validation loss to improve. Churn tracks the best loss seen so far; every epoch without a new best adds one to a counter; at 25 the loop stops and the best state_dict — the saved weights, also called a checkpoint — is the result, not the last one. ReduceLROnPlateau is a second, slower waiter: it watches the same loss and halves the learning rate after 10 stalled epochs.

1.7 · Accuracy, precision, recall

Accuracy is the fraction of rows whose predicted class matches the label: after sigmoid above 0.5 for churn, after argmax (the index of the largest logit) over the 10 logits for FashionMNIST. It hides imbalance. The churn table is 73.5% no-churn (5,174 of 7,043), so a model that always says "no" scores 73.5% without learning anything; the run's 0.76 is about 2.5 points above that floor. Precision is \(\mathrm{TP}/(\mathrm{TP} + \mathrm{FP})\): of the customers flagged, how many churned. Recall is \(\mathrm{TP}/(\mathrm{TP} + \mathrm{FN})\): of the customers who churned, how many were flagged. A retention list at a fixed budget wants precision; a safety check wants recall.

2 · Churn: the binary problem

A telecom table, one row per customer, binary target: churned next month or not. 1,869 of the 7,043 rows churn, 26.5%. Acquisition costs 500–1000 per new customer against 50–200 monthly revenue per churned one; a retention campaign costs 50–100 per targeted customer at roughly 60% success. So the model is a triage list for who to target, and precision on the minority class is the point. The pandas cleaning and EDA walkthrough are assumed, so this page starts where torch starts.

2.1 · Package imports

torch.nn as nn gives layers (Linear, BatchNorm, Dropout, ReLU). torch.optim gives Adam/SGD; torch.optim.lr_scheduler lowers the learning rate when validation loss stops improving. sklearn.metrics gives the report card (precision, recall, and a confusion matrix: predicted class against true class). train_test_split with stratify holds the class ratio across splits.

3 · Churn architecture: 40 in, 256-128-64, BatchNorm

ChurnPredictor(nn.Module) builds its hidden stack in a loop: Linear → BatchNorm → ReLU → Dropout, widths 256 → 128 → 64, dropout 0.2, then one Linear to a single logit. The first Linear reads 40 numbers: the raw table has 21 columns; dropping the ID and target, adding 11 engineered columns, and one-hot encoding the four multi-category columns (one 0/1 column per category) leave 40. BatchNorm normalizes the input to every hidden layer with learned parameters; dropout is regularization for later weeks. The forward method just feeds x through self.net. super().__init__() is load-bearing: without it layers never register, parameters() is empty, and nothing trains.

class ChurnPredictor(nn.Module):
    def __init__(self, input_size, hidden=(256, 128, 64), p=0.2):
        super().__init__()  # registers layers; without it, parameters() is empty
        layers, prev = [], input_size
        for h in hidden:                       # 256 -> 128 -> 64
            layers += [nn.Linear(prev, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(p)]
            prev = h
        layers += [nn.Linear(prev, 1)]         # single logit, no sigmoid here
        self.net = nn.Sequential(*layers)
    def forward(self, x):
        return self.net(x).squeeze(1)          # (N,1) -> (N,) for the loss

One logit out, sigmoid lives inside the loss. Squeezing to (N,) matches BCEWithLogitsLoss against a (N,) target.

The stack holds 52,609 parameters. The first Linear has \(40 \times 256 + 256 = 10{,}496\), the second \(256 \times 128 + 128 = 32{,}896\), the third \(128 \times 64 + 64 = 8{,}256\), and the head \(64 \times 1 + 1 = 65\). The three BatchNorm layers add two parameters per channel, \(2(256 + 128 + 64) = 896\). The sum is \(10{,}496 + 32{,}896 + 8{,}256 + 65 + 896 = 52{,}609\).

The 40 is the feature count, not the batch size. Change it and the first matrix multiply stops.

4 · Init: Xavier on linears, zeros on biases

PyTorch defaults already initialize, but here the linears are overridden with Xavier. Loop model.modules(), keep only nn.Linear, apply Xavier to weights, zeros to biases. BatchNorm and Dropout are skipped — they need no Xavier. Zero bias avoids an arbitrary offset at the start.

for m in model.modules():
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)  # variance 2/(n_in + n_out)
        nn.init.constant_(m.bias, 0.0)     # no offset at the start

The uniform bound is sqrt(6/(n_in+n_out)), whose variance matches the Xavier target. BatchNorm and Dropout hold no weight matrix, so they are skipped.

5 · Split, loss, optimizer

First hold out 20% of the 7,043 rows for test: 1,409 rows. Then split the remaining 5,634 into 4,508 train and 1,126 validation, again 80/20 and again stratified, so both pieces keep the 26.5% churn ratio. Convert back through NumPy because the splitter takes NumPy, then return to torch on device. The loss is BCEWithLogitsLoss(pos_weight=1.88): sigmoid plus binary cross-entropy in one stable step, with the positive weight taken from the training counts. Adam starts at 1e-4 with weight decay 1e-5; ReduceLROnPlateau halves the rate after 10 epochs without a validation-loss improvement.

from sklearn.model_selection import train_test_split

Xtr_full, Xte, ytr_full, yte = train_test_split(X.cpu().numpy(), y.cpu().numpy(),
    test_size=0.20, stratify=y.cpu().numpy())        # 5,634 / 1,409
Xtr, Xva, ytr, yva = train_test_split(Xtr_full, ytr_full,
    test_size=0.20, stratify=ytr_full)               # 4,508 / 1,126
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)  # imbalance lives here
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
sched = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min",
    factor=0.5, patience=10)

The first call holds out test; the second holds out validation from what remains. The scaler is fitted on the 5,634 rows before the second split, so validation rows enter the scaler's mean and standard deviation. Fit it on the 4,508 training rows only; the leak is small here, but the rule has no exception.

The status line prints lr=0.001, but the optimizer call sets 1e-4 and the LR column shows 0.000100; the call is what runs.

6 · The churn train loop: 500 epochs, patience 25

One epoch is one train pass plus one validation pass. Train: model.train(), forward on the full split, loss, optimizer.zero_grad() (grads accumulate otherwise), loss.backward(), optimizer.step(). Validate: model.eval(), forward only, no optimizer, record loss and accuracy. Step the scheduler on validation loss. Track best validation loss from infinity, patience counter to 25, save the best state_dict. Print every 10 epochs: train loss, validation loss, validation accuracy, learning rate. On the stored run the loop stopped at epoch 492, not 500: the best validation loss was 0.6064 and 25 epochs passed without a new best. The printed rate was 1e-4 at epoch 471, 5e-5 at 481, and 2.5e-5 at 491, as the scheduler halved it.

best, wait = float("inf"), 0
for epoch in range(500):
    model.train()
    pred = model(Xtr); loss = criterion(pred, ytr)
    optimizer.zero_grad(); loss.backward(); optimizer.step()  # the relay in torch
    model.eval()
    with torch.no_grad():                                     # validation learns nothing
        vp = model(Xva); vloss = criterion(vp, yva)
    sched.step(vloss)
    if vloss < best: best, wait = vloss, 0; torch.save(model.state_dict(), "churn_best.pth")
    else: wait += 1
    if wait >= 25: break                                     # early stopping

model.load_state_dict(torch.load("churn_best.pth"))          # the best weights, not the last

The three verbs in order — zero, backward, step — are one update. The order is the contract.

The test split is 1,409 rows, 374 of them churners. Accuracy 0.76 sits 2.5 points above the always-no floor of 0.735.

7 · FashionMNIST: the 10-way setup

60,000 train and 10,000 test grayscale images, \(28 \times 28 \times 1\), pixels 0–255, 10 clothing classes. Same format as MNIST, harder shapes — the first step toward vision, still an MLP. The loader uses batch size 64: 938 batches cover the training set and 157 cover the test set. Architecture flattens 784 → 128 → 64 → 10 logits; softmax turns a row of logits into a distribution, and the loss applies it internally. Hidden layers funnel detail into abstraction with ReLU (gradient 1 when open, which keeps the signal alive); the output is a distribution over 10 classes, not one threshold at 0.5.

Preprocessing composes ToTensor (0–255 → 0–1) with Normalize at the dataset mean and std. Forward takes \((B, 1, 28, 28)\), flattens to \((B, 784)\), runs Linear → ReLU → Dropout twice, then Linear to \((B, 10)\) logits. No BatchNorm here — left as your exercise. A torch.no_grad() test pass on a random \((32, 1, 28, 28)\) batch should return \((32, 10)\); softmax turns those logits into probabilities.

The stack holds 109,386 parameters: \(784 \times 128 + 128 = 100{,}480\), \(128 \times 64 + 64 = 8{,}256\), and \(64 \times 10 + 10 = 650\).

self.flatten = nn.Flatten()          # (B,1,28,28) -> (B,784)
x = self.flatten(x)                  # 784 = 28*28
x = self.drop(torch.relu(self.fc1(x)))  # 784 -> 128
x = self.drop(torch.relu(self.fc2(x)))  # 128 -> 64
logits = self.fc3(x)                 # 64 -> 10, softmax lives in the loss

8 · Train one epoch, validate without grad

Loss is CrossEntropyLoss (log-softmax plus negative log-likelihood, NLL, in one step — never softmax first). Adam at lr 0.001 with weight decay (L2 penalty keeping weights small). Accuracy helper: argmax(dim=1) picks the top logit per sample, compare with labels, mean the hits. Train epoch: model.train() (dropout on), 938 minibatches of 64, per-minibatch zero_grad → forward → loss → backward → step, accumulate loss and accuracy, print every 200 batches. Validation helper evaluate: model.eval() (dropout off), torch.no_grad() (faster, less memory), accumulate loss/accuracy plus all predictions and labels for confusion matrices later.

model.train()
for xb, yb in train_loader:          # one minibatch at a time
    xb, yb = xb.to(device), yb.to(device)
    optimizer.zero_grad()
    loss = criterion(model(xb), yb)
    loss.backward(); optimizer.step()

model.eval()
with torch.no_grad():                # evaluation learns nothing
    for xb, yb in test_loader:
        loss = criterion(model(xb.to(device)), yb.to(device))

Train mode enables dropout/BatchNorm updates; eval mode disables them. Forgetting model.eval() is a silent validation bug.

9 · The outer loop and the overfit gap

Overfitting is a training score that keeps improving while the validation score stalls; the gap between the two measures it. Outer loop over epochs: record start time, run train_epoch, run evaluate, append all four curves, print train/test loss plus epoch time. Optional stop past a target, say test accuracy above 97%. Then plot train vs test loss, train vs test accuracy, and the accuracy improvement; report best and final accuracies plus the gap train − test. Above 10 points is significant overfitting, above 5 is mild. That gap is what §6's patience-25 rule was defending against on churn.

The run used 10 epochs, Adam at 1e-3 with weight decay 1e-4. Epoch 1: train loss 0.5482, test loss 0.4211, train accuracy 80.12%, test 84.32%. Epoch 10: train loss 0.2982, test loss 0.3278, train accuracy 88.99%, test 87.99%. The test loss starts lower in epoch 1 because dropout is on during training and off during evaluation; the two losses are not measured under the same conditions. The gap is 1.00 point, below the 5-point mild line, so this MLP generalizes on this split.

The stop rule reads test accuracy every epoch, so the 10,000 test images act as a validation set during the run. The final 87.99% is the last epoch's score on those images, but the choice to stop was made with the same images. A clean split would choose on a validation slice and read test once.

10 · The code, and what happens when …

The contracts above are easy to read and easy to break. Each experiment asks for a prediction before the code. Write the prediction down; the paragraph after the code is the check.

10.1 · What happens with the wrong feature count?

Predict first. The table has 21 columns. The model's first layer is Linear(40, 256). You pass a batch of 4 rows that still has the 21 raw columns. Does it run, and if not, what does the message name?

import torch, torch.nn as nn

first = nn.Linear(40, 256)          # 40 engineered features in
x = torch.zeros(4, 21)              # 4 rows, 21 raw columns
print(first(x))

What happens. RuntimeError: mat1 and mat2 shapes cannot be multiplied (4x21 and 40x256). A Linear layer is a matrix multiply, and the inner dimensions 21 and 40 do not match, so PyTorch stops before any arithmetic. It does not pad, truncate, or broadcast the feature axis; broadcast only stretches a dimension of size 1. The 40 is the feature count after the 21-column table loses its ID and target, gains 11 engineered columns, and one-hot encodes four categorical columns; the 4 is the batch size and has nothing to do with it.

The rule: the first Linear's in_features must equal the tensor's last dimension. Build the model with input_size=40, not 21.

10.2 · What happens when you softmax before the loss?

Predict first. Three rows, three classes, targets 0, 1, 2. The first loss gets the raw logits. The second gets the same numbers after softmax. Which loss is larger, and by how much?

logits = torch.tensor([[2.0, 1.0, 0.0],
                       [0.0, 2.0, 1.0],
                       [1.0, 1.0, 1.0]])
target = torch.tensor([0, 1, 2])
ce = nn.CrossEntropyLoss()

print(round(ce(logits, target).item(), 4))             # 0.6379
print(round(ce(logits.softmax(1), target).item(), 4))  # 0.8977

What happens. The second loss is larger: 0.8977 against 0.6379. CrossEntropyLoss applies log-softmax inside, so the second call computes log-softmax of probabilities, which pulls the row toward uniform and lowers the probability of the true class. The gradient shrinks the same way. Pass raw logits; apply softmax only when you want to read probabilities.

The rule: one softmax per row, and it lives inside the loss. A model that ends in softmax has applied it twice by the time the loss runs.

10.3 · What happens to BatchNorm in eval mode?

Predict first. A fresh BatchNorm1d(1) sees 1, 2, 3, 4 in train mode. Its four outputs have mean 0. The same module then switches to eval mode and sees the same four numbers. Are the outputs still centred on 0?

bn = nn.BatchNorm1d(1)
x = torch.tensor([[1.], [2.], [3.], [4.]])

bn.train()
print([round(v, 4) for v in bn(x).flatten().tolist()])
# [-1.3416, -0.4472, 0.4472, 1.3416]

bn.eval()
print([round(v, 4) for v in bn(x).flatten().tolist()])
# [0.7262, 1.6944, 2.6627, 3.6309]
print(round(bn.running_mean.item(), 4), round(bn.running_var.item(), 4))
# 0.25 1.0667

What happens. Train mode subtracts the batch mean 2.5 and divides by the batch standard deviation 1.118, so the four outputs have mean 0. Eval mode uses the running statistics that the train call just updated: mean 0.25, variance 1.0667. The output is shifted and scaled, not re-centred, and a row's value no longer depends on the other three rows. That is the point of eval: the same row must get the same score in every batch. A validation pass that forgets model.eval() measures a moving target.

The running variance uses the unbiased estimate (divide by \(n-1\)): \(0.9 \times 1 + 0.1 \times 1.6667 = 1.0667\).

Hands-on

1. Write the four churn shapes before you run the split: the full table, the test holdout, the training split, the validation split. Predict 7,043 × 40, 1,409 × 40, 4,508 × 40, 1,126 × 40, then print them. Check the two ratios stay 80/20 and both classes appear on every side.
2. Count the first Linear by hand: 40 inputs × 256 outputs + 256 biases. Predict 10,496, then sum p.numel() over model.net[0].parameters(). Add the other layers and check 52,609.
3. Feed 21 columns to model. Predict the error text before you run it. Read the two shapes in the message, then pad the batch to 40 columns with zeros and check that the forward pass runs.
4. Run BatchNorm1d(1) on 1, 2, 3, 4 in train mode, then in eval mode. Predict the mean of each output before you print. Train gives 0; eval gives 2.18. Say which statistics each call used.
5. Compute CrossEntropyLoss on the 3×3 logits and on their softmax. Predict which is larger. Check 0.6379 and 0.8977, then delete the softmax and keep the logits.
6. Compute the always-no accuracy, 5,174/7,043. Predict the gap to the model's 0.76. Read the test report and write one sentence on what churn precision 0.54 adds that accuracy 0.76 hides.

One-line understanding. Build the stack, initialize it so no two hidden units start equal, split without leaking, then loop zero-backward-step on train and no-grad forward on validation until patience runs out.