Module 10 — A first network trained end to end

Nine modules of mechanisms. This module assembles them into a complete, runnable project: a classifier on tabular data, from preprocessing to saving. Every decision in the code points back to the module that justifies it.

The project's structure

The order of steps is not negotiable, and it follows directly from the previous course: split before preprocessing, fit the preprocessing on training data only, then build the network.

import torch
import torch.nn as nn
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# 1. Split first, preprocess after: the rule from the previous course.
X_tr, X_val, y_tr, y_val = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42
)

# 2. The scaler is fitted on training data only, otherwise there is leakage.
scaler = StandardScaler().fit(X_tr)
X_tr = scaler.transform(X_tr)
X_val = scaler.transform(X_val)

# 3. Tensors then loaders. Shuffle training, never validation.
def to_tensors(X, y):
    return torch.utils.data.TensorDataset(
        torch.tensor(X, dtype=torch.float32),
        torch.tensor(y, dtype=torch.long),
    )

loader_tr = torch.utils.data.DataLoader(to_tensors(X_tr, y_tr), batch_size=64, shuffle=True)
loader_val = torch.utils.data.DataLoader(to_tensors(X_val, y_val), batch_size=256)

The network

class Classifier(nn.Module):
    def __init__(self, n_inputs, n_classes, width=128, dropout=0.3):
        super().__init__()
        self.network = nn.Sequential(
            # No bias: the normalization's beta parameter replaces it.
            nn.Linear(n_inputs, width, bias=False),
            nn.BatchNorm1d(width),
            nn.ReLU(),
            nn.Dropout(dropout),

            nn.Linear(width, width // 2, bias=False),
            nn.BatchNorm1d(width // 2),
            nn.ReLU(),
            nn.Dropout(dropout),

            # Raw scores out: the loss applies the softmax itself.
            nn.Linear(width // 2, n_classes),
        )

    def forward(self, x):
        return self.network(x)


device = "cuda" if torch.cuda.is_available() else "cpu"
model = Classifier(X_tr.shape[1], len(np.unique(y))).to(device)

Three decisions in this block deserve to be spelled out. The order linear, normalization, activation, dropout follows the convention of module 8. The bias=False is explained by the normalization's $\beta$ parameter, which already fills that role. And the last layer carries no activation: CrossEntropyLoss expects raw scores and applies the softmax in a numerically stable way, as module 3 explained.

The training loop

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

best_loss, patience, waited = float("inf"), 10, 0
history = {"train": [], "val": []}

for epoch in range(100):
    # --- Training: dropout active, batch statistics. ---
    model.train()
    loss_tr = 0.0
    for X_batch, y_batch in loader_tr:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)
        optimizer.zero_grad()                  # otherwise gradients accumulate
        loss = criterion(model(X_batch), y_batch)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # module 6
        optimizer.step()
        loss_tr += loss.item() * X_batch.size(0)

    # --- Validation: dropout disabled, running averages. ---
    model.eval()
    loss_val, correct = 0.0, 0
    with torch.no_grad():
        for X_batch, y_batch in loader_val:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            outputs = model(X_batch)
            loss_val += criterion(outputs, y_batch).item() * X_batch.size(0)
            correct += (outputs.argmax(1) == y_batch).sum().item()

    loss_tr /= len(loader_tr.dataset)
    loss_val /= len(loader_val.dataset)
    history["train"].append(loss_tr)
    history["val"].append(loss_val)
    scheduler.step()

    # --- Early stopping: keep the weights of the best pass. ---
    if loss_val < best_loss:
        best_loss, waited = loss_val, 0
        torch.save(model.state_dict(), "best_model.pt")
    else:
        waited += 1
        if waited >= patience:
            print(f"early stopping at epoch {epoch}")
            break

    if epoch % 10 == 0:
        accuracy = correct / len(loader_val.dataset)
        print(f"epoch {epoch:3d} | train {loss_tr:.4f} | val {loss_val:.4f} | accuracy {accuracy:.3f}")

Four points of vigilance concentrated in this loop, all drawn from previous modules. Switching between train() and eval() governs the behavior of dropout and normalization (modules 7 and 8). The zero_grad() prevents gradient accumulation across iterations. The torch.no_grad() in validation saves memory and compute by not recording the graph. And the scheduler advances once per epoch, not per batch.

Verify, then save

Before launching those hundred epochs, apply the test from module 9: overfit ten observations. If it fails, fix that before going further.

For deployment, reload the best weights and save everything needed to reproduce the result:

model.load_state_dict(torch.load("best_model.pt"))
model.eval()

torch.save({
    "weights": model.state_dict(),
    "architecture": {"n_inputs": X_tr.shape[1], "width": 128, "dropout": 0.3},
    "scaler": {"mean": scaler.mean_, "std": scaler.scale_},
    "torch_version": torch.__version__,
}, "full_model.pt")

Saving the scaler alongside the weights is indispensable, and a frequent omission. Without the training mean and standard deviation, the model will receive production data at a different scale from the one it learned, and its predictions will be meaningless. This is exactly the pipeline logic of the previous course, transposed to PyTorch.

A realistic path of improvement

This code is a solid foundation, but the order of improvements matters. Look first for the right learning rate, which pays more than everything else. Then width and depth, watching the gap between the two curves. Then the dropout rate and the weight penalty together. Only then add architectural complexity. And keep module 1's point in mind: on tabular data, a well-tuned gradient boosting will often beat this network — the question to ask before starting is about the nature of the data.

Summary

• Split before preprocessing, and fit the scaler on the training set only: the previous course's rule applies unchanged.
• The order linear, normalization, activation, dropout, with bias=False before a normalization and no activation on the output.
• The loop requires train() and eval() in the right places, zero_grad(), no_grad() in validation, and a scheduler advanced per epoch.
• Save the weights and the scaler together; without the scaling statistics, the model is unusable in production.

Final step: the recap and the 40-question exam.