Building and Training an MLP You Can Defend

You have already subclassed nn.Module, defined forward, and seen that the code works. The important question now is: what does nn.Module actually keep track of for us?

Start with the concrete idea. A neural network contains numerical values that training is allowed to change. Those values are the model’s parameters — for example, the weights and bias of a Linear layer.

PyTorch needs to know which objects inside your Python object are part of the neural network. When you assign a layer such as nn.Linear(...) to an attribute of an nn.Module inside __init__, PyTorch records that layer as a child module. The parameters inside that layer are then part of the model’s registered parameters.

This tracking is called parameter/module registration. It is not cosmetic bookkeeping. Later, several PyTorch operations rely on this record:

• model.parameters() finds the trainable tensors that an optimizer should update;
• model.to(device) moves registered parameters and buffers to CPU, CUDA, or another device;
• state_dict() collects the registered model state so it can be saved;
• model.train() and model.eval() propagate the appropriate mode to child modules.

So when this lesson says that a layer is “registered”, read that as: PyTorch knows that this object belongs to the model and can include it in these model-management operations.

An optimizer is the object that changes model parameters after gradients have been calculated. Adam is one example.

When you write:

torch.optim.Adam(model.parameters(), ...)

model.parameters() gives Adam the tensors it is allowed to update. This is why registration matters: if a layer’s parameters are not registered, the optimizer never receives them.

The dangerous part is that this usually does not produce an error. The model can still run a forward pass, calculate a loss, and run backward(). It simply leaves some weights unchanged. That is a silent training bug: the program runs, but the model is only partly learning.

That is the same category of bug as Lesson 2.1’s silent broadcast: the program runs, produces plausible numbers, and does not necessarily tell you that something is wrong. The defence is the same — check the structure instead of assuming it:

assert count_params(model) == expected, "parameter count mismatch"

If you need a collection of layers, use a PyTorch container that knows it is holding modules:

• nn.ModuleList is a list-like container for child modules;
• nn.ModuleDict is a dictionary-like container for named child modules.

A plain Python list or dictionary does not automatically register the modules it contains.

There is another layer of machinery around forward. When you write:

model(x)

Python invokes the module’s __call__ machinery, which in turn handles the forward pass and PyTorch features such as hooks. Calling model.forward(x) directly jumps past that machinery.

So the normal rule is simple: call the model as model(x), not model.forward(x). You define forward; PyTorch’s module interface decides how that forward computation is wrapped.

A network is easier to reason about when you can describe it as a sequence of reusable pieces rather than as one long expression.

There are three common ways to assemble a network, and they trade brevity against control:

1. nn.Sequential — useful when the data simply flows through a fixed sequence of layers;
2. a reusable sub-module — useful when a pattern such as Linear → ReLU appears more than once;
3. a custom nn.Module with named child modules — useful when the forward pass contains branching, skip/residual connections, multiple inputs, or other logic that a straight sequence cannot express.

The important idea is composition: a large model can be built from smaller modules, and those modules can themselves contain modules.

Real neural-network architectures often repeat the same small pattern. For this model the repeated pattern is essentially:

Linear → ReLU

Instead of writing that pattern repeatedly, we can package it into a small module called a block. A block is simply a reusable piece of a larger network.

The technical idea is composition: we construct a larger model by putting smaller modules together. This makes the model easier to read, inspect, change, and reuse. The same design idea appears in much larger architectures, including the repeated blocks used in Transformers.

The dotted parameter names are not arbitrary labels. They describe the path through the module tree.

A module tree is simply the nested structure of the model. For example:

BankMLP
├── blocks
│   ├── 0
│   │   ├── linear
│   │   └── act
│   └── 1
│       ├── linear
│       └── act
└── head

So blocks.0.linear.weight means:

BankMLP → blocks → block 0 → linear layer → its weight tensor.

Those paths become keys in state_dict(), which is PyTorch’s dictionary-like representation of the model state. This is why renaming a submodule can affect whether an old checkpoint loads cleanly: the saved keys may no longer match the new module-tree paths.

This is also why Section 8 saves the architecture configuration alongside the weights. The weights are meaningful only when we reconstruct the architecture they belong to.

Lesson 2.1’s shape clinic dealt with one layer at a time. In a real model, a shape mistake may originate several layers before the line that finally fails. The error message usually identifies the operation that could not proceed, not necessarily the earlier configuration choice that caused the mismatch.

The first tool is deliberately simple: dry-run a fake batch before you touch real data.

A dry run means sending artificial data through the model only to test the wiring. We do not care about the predictions. We care that the input shape is accepted and that the output shape is what the next stage expects.

A hook is a function that PyTorch can call automatically when a particular event happens inside a module. A forward hook runs when a module has performed its forward computation, so it lets us inspect what went into a layer and what came out without rewriting the model.

Think of a forward hook as temporarily attaching a measuring instrument to a layer.

Register one on each leaf submodule and run one fake batch. You then get a trace of the tensor as it moves through the network.

This also explains why the earlier distinction between model(x) and model.forward(x) matters. Hooks are part of the module’s __call__ machinery. Calling forward directly bypasses that machinery, so a forward hook would not fire in the normal way.

Read the output column downward: 50 → 64 → 64 → 32 → 32 → 1.

Here is how to read that trace:

• the first Linear changes the feature width from 50 to 64;
• ReLU changes the values but not their shape, so (4, 64) stays (4, 64);
• the second Linear changes 64 features to 32;
• the second ReLU leaves (4, 32) unchanged;
• the final Linear produces one raw output per example, (4, 1).

The first dimension, 4, is the batch dimension. It is the number of examples processed together. None of these layers changes that dimension.

Hooks must also be removed after debugging. A registered hook remains attached to the module. If it keeps appending records to a Python list, that list can continue growing every time the model runs. Leaving diagnostic hooks installed can therefore create unnecessary memory use and confusing later measurements.

Here is a network with a deliberate mismatch. The important debugging idea is that the traceback names the operation that failed. A shape trace lets you work backwards from the failure and see the last successful shape.

This is a binary classification problem: each customer is either in the negative class (0) or the positive class (1). There are two standard ways to represent that binary choice in the output layer.

Before the table, learn the vocabulary:

• A logit is a model’s raw numerical score before it has been converted into a probability. A large positive logit favours the positive class; a large negative logit favours the negative class; a value near zero represents uncertainty.
• A probability is a number between 0 and 1 that we can interpret as the model’s estimated chance of the positive class.
• Sigmoid converts one logit into one number between 0 and 1.
• Softmax converts a vector of class scores into values that sum to 1, so they can be interpreted as probabilities across classes.

With those ideas in place, the two formulations are:

Both formulations are valid for two classes. The two-logit formulation is convenient when the same code will later be extended to more than two classes. The one-logit formulation is leaner for this binary task and gives us the pos_weight argument we will use later to demonstrate class weighting.

This lesson uses one logit. That is why Lesson 2.2 prepared the labels as (N, 1) floating-point values.

BCEWithLogitsLoss may look like a convenience wrapper around sigmoid followed by BCELoss, but the important difference is numerical stability.

First, the terms:

• Binary cross-entropy (BCE) is a loss function for a binary target. It penalises predictions that assign low probability to the correct class.
• Sigmoid converts a raw score into a probability.
• Numerical stability means arranging the arithmetic so that a computer can calculate it reliably even for very large or very small numbers. A mathematically equivalent formula can behave differently in floating-point arithmetic.

Why does this matter? During training a model can produce very large positive or negative logits. If we first apply sigmoid, those logits can be pushed extremely close to 1 or 0. A direct probability-based calculation can then encounter values that are effectively rounded or clamped at the limits of floating-point precision. The corresponding gradient can become unhelpfully small or behave poorly.

BCEWithLogitsLoss combines the sigmoid and BCE calculations into one formulation designed to avoid those numerical problems. Therefore the model should normally output the raw logit, and the loss should receive that raw logit.

The cell below compares the stable formulation with the naive “sigmoid first” route, especially at logits where the model is confidently wrong.

The rule that follows:

During training, do not put a sigmoid or softmax in this model’s forward when the corresponding loss already expects logits. Output raw logits and let the loss function perform the appropriate transformation internally. Apply sigmoid later when you explicitly need a probability for interpretation, ranking, thresholding, or a metric.

If you put a sigmoid in forward and then use BCEWithLogitsLoss, the loss receives a value that has already been squashed into the range 0–1 even though it expects an unconstrained logit. You have effectively given the loss the wrong quantity. The code can still run, so this is another silent bug: the problem is in the mathematics, not in Python syntax.

That comparison is a useful sanity check on a new classification problem: calculate what a deliberately simple model would score before judging a trained model.

Here, predicting roughly the same probability for everyone gives us a reference loss. This is a baseline, not a mathematical lower bound: it tells us what happens with a simple strategy, so we have something meaningful to compare against from the first epoch.

The six-line mechanism from Module 1 has not changed:

optimizer.zero_grad()             # clear gradients from the previous step
predictions = model(xb)           # forward pass: compute the model's output
loss = loss_fn(predictions, yb)   # measure how wrong that output is
loss.backward()                   # calculate gradients of the loss w.r.t. parameters
optimizer.step()                  # use those gradients to update the parameters

To understand the loop, connect each line to the question it answers:

1. What should the model forget from the previous update? → clear the old gradients.
2. What does the model currently predict? → run the forward pass.
3. How wrong is that prediction? → calculate the loss.
4. Which parameters contributed to that error, and in what direction should they move? → backpropagation calculates gradients.
5. What should actually change? → the optimizer updates the registered parameters.

A real training loop adds relatively little new mathematics. Most of the extra code is bookkeeping: keeping correct totals, switching modes, moving tensors to the right device, and recording metrics without accidentally keeping computation graphs alive.

1. Forgetting zero_grad().

A parameter’s gradient is stored in the parameter object, and PyTorch accumulates new gradients into the existing value by default. If you do not clear the previous gradient before the next batch, the optimizer uses a mixture of information from multiple batches. That is usually not the update you intended.

2. Accumulating loss instead of loss.item().

loss is a tensor that still belongs to the computation graph — the chain of operations PyTorch recorded so it can calculate derivatives during backward(). If you repeatedly add those tensors to a running total, you can keep references to all of those graphs instead of just keeping the numerical loss value.

loss.item() extracts the ordinary Python number from a single-element tensor. Once you have that number, the old computation graph is not kept alive by your running total.

3. Averaging the averages.

loss_fn returns the mean loss for the current batch. If every batch had exactly the same number of examples, averaging those batch means would also give the dataset mean. But the final batch is often shorter.

Suppose two batches contain 256 examples and the final batch contains 10. Giving each batch loss equal weight makes the 10-example batch count as much as a 256-example batch. The correct dataset mean weights each batch by its actual number of examples:

sum of per-example losses across all examples
----------------------------------------------
number of examples

That is why the code multiplies each batch mean by xb.size(0) and divides the total by the total number of examples.

The validation pass looks like the training pass with three lines removed. Those three absences are the entire point.

These two lines often appear together, which can make them look like two spellings of the same idea. They control completely different mechanisms.

model.eval() controls how certain layers behave.

A model can have layers whose behaviour changes between training and inference. Dropout, for example, randomly removes activations during training to regularise the model; during evaluation it should not randomly remove them. BatchNorm also changes how it obtains its statistics. Calling model.eval() tells those modules to use their evaluation behaviour. The setting is propagated through the model’s child modules.

eval() does not turn gradients off. Parameters can still participate in gradient tracking if the surrounding code allows it.

torch.no_grad() controls whether PyTorch builds the gradient graph.

During validation we normally do not need to calculate gradients because we are measuring the model, not updating it. torch.no_grad() tells PyTorch not to record the operations needed for a later backward pass. That reduces memory use and usually makes evaluation faster.

So remember the distinction:

model.eval()       → change layer behaviour
with torch.no_grad() → stop building the gradient graph

BankMLP currently contains neither dropout nor batch normalization, so eval() does not visibly change this particular model. We still use it because the correct habit matters once the architecture contains layers whose behaviour depends on training versus evaluation mode.

model.train() and model.eval() are state changes. Whichever one ran most recently remains in effect until another call changes it.

That means a robust training function should set model.train() when training begins, and a robust evaluation function should set model.eval() when evaluation begins. Then the functions do not depend on the order in which you called them earlier.

A mode leak is what happens when a function accidentally inherits the mode left behind by another part of the program. For example, if evaluation leaves the model in evaluation mode and the next training epoch forgets to call train(), a model containing dropout would train with dropout disabled for that epoch.

This kind of bug can be difficult to notice because the program still runs and the numbers may look plausible.

Before interpreting a model’s score, establish what simple alternatives can already achieve. A model’s metric is meaningful only relative to an appropriate baseline.

Two baselines are used here:

• Majority-class baseline: predict the most common class for everyone. It is useful for showing why accuracy is misleading on an imbalanced dataset. If 88.3% of customers are negative, always predicting negative gives about 88.3% accuracy while identifying none of the positive customers.
• Logistic regression: a linear classifier trained on the same encoded features. This is the more important modelling baseline because it asks a concrete question: does the extra nonlinear machinery of the MLP provide enough additional predictive value to justify its additional complexity and training cost?

The baseline is not something to “beat” merely because it is simple. It is the reference that makes the result interpretable.