Building MLPs in PyTorch

NumPy provides fast numerical arrays and vectorized math. All matrix multiplications, element-wise operations, and random number generation in this code run as optimized C code under the hood.

The two-moons dataset is a classic 2D classification benchmark. It consists of two interleaving crescent shapes that cannot be separated by a straight line — requiring a non-linear classifier like our MLP. Unlike the diagonal flip (which was a regression task), this is binary classification: each point belongs to class 0 or class 1.

Fixes the random number generator so we get identical data every run. This is critical for fair comparisons — the only variable should be the architecture, not random luck in data generation.

We generate 150 points for each class (300 total). This is large enough to learn a good boundary but small enough to train quickly with full-batch gradient descent.

Creates 150 evenly spaced angles from 0 to π (a half-circle). These angles trace out the crescent shapes. linspace(0, π, 150) = [0, 0.0211, 0.0422, ..., 3.1416].

Creates the upper crescent. cos(t) gives x-coordinates, sin(t) gives y-coordinates. As t goes from 0 to π, this traces the top half of a unit circle: starting at (1, 0), arcing through (0, 1), ending at (-1, 0).

Creates the lower crescent. Same semicircle as the upper moon, but flipped vertically (-sin) and shifted right (+0.5) and up (+0.5). The offset makes the two moons interleave — they overlap in the middle, creating a non-linearly-separable pattern.

Stacks upper and lower moon into one dataset. Rows 0–149 are class 0 (upper), rows 150–299 are class 1 (lower).

Adds Gaussian noise to make the task realistic. Without noise, the moons are clean curves that could be perfectly separated. With noise = 0.1, points scatter slightly around the crescents, creating overlap near the boundary. This forces the network to learn a robust decision boundary.

Creates binary labels: 0 for upper moon (first 150 points), 1 for lower moon (last 150 points). These are the targets our classifier will learn to predict.

Creates a random ordering of indices 0–299. We use this to shuffle the data so class 0 and class 1 are interleaved, not grouped. This prevents the network from learning spurious patterns based on sample order.

Reorders both X and y using the same permutation. This ensures each data point keeps its correct label after shuffling.

Confirms the dataset dimensions. Output: Dataset: X(300, 2), y(300,). 300 points, 2 features per point, 1 label per point.

The sigmoid function squashes any real number into the range (0, 1), making it interpretable as a probability. For classification, the output neuron uses sigmoid (not ReLU) so we get P(class=1|x). Sigmoid is defined as σ(z) = 1/(1 + e^(-z)).

Clips extreme values to prevent numerical overflow in exp(-z). Without clipping: exp(-(-1000)) = exp(1000) = infinity. With clipping: exp(-(-500)) = exp(500) is still huge but finite. The sigmoid of any z > 20 is effectively 1.0, so clipping at 500 has zero effect on the result.

The core sigmoid computation: σ(z) = 1/(1 + e^(-z)). When z is large and positive, e^(-z) ≈ 0, so σ ≈ 1. When z is large and negative, e^(-z) is huge, so σ ≈ 0. When z = 0, e^0 = 1, so σ = 1/2 = 0.5.

We build a 2-input, 16-hidden, 1-output network. 2 inputs because our data is 2D (x, y coordinates). 16 hidden neurons with ReLU (from Section 1, we know 16 neurons give ample capacity). 1 output with sigmoid for binary probability.

Sets a separate seed for weight initialization. Using seed 0 (different from data seed 42) ensures independence between data randomness and weight randomness.

Creates the weight matrix for the hidden layer using He initialization. Shape (2, 16) means 2 inputs and 16 hidden neurons. He scale: √(2/fan_in) = √(2/2) = 1.0, so we just use standard normal values.

One bias per hidden neuron, all initialized to zero. Standard practice — He initialization handles the variance through weights.

Weight matrix for the output layer. Shape (16, 1): 16 inputs from hidden layer, 1 output. He scale: √(2/16) = √0.125 = 0.354.

Single output bias. With b2 = 0 and random weights, the initial sigmoid output is centered around 0.5 (no class preference).

Learning rate for gradient descent. We use lr = 1.0 because the gradient is averaged over 300 samples (divided by N), making the effective step size reasonable. With lr = 0.1, convergence would be 10× slower.

Train for 500 epochs. Each epoch does one full forward pass (all 300 samples), computes the loss, does one full backward pass, and updates all weights. This is full-batch gradient descent — every epoch uses the entire dataset.

The forward pass computes predictions for all 300 samples simultaneously using matrix operations. Data flows: input (300,2) → hidden (300,16) → output (300,1) → probability (300,).

Computes the weighted sum for every hidden neuron on every sample. X is (300, 2) and W1 is (2, 16), so X @ W1 is (300, 16) — 300 samples, 16 pre-activations each. Adding b1 broadcasts the (16,) bias across all 300 rows.

Applies ReLU element-wise: negative values become 0, positive values pass through. This creates the non-linearity that lets the MLP learn curved boundaries. ~50% of values are zeroed out (sparse activation).

Computes the raw output (logit) by multiplying hidden activations by output weights. H is (300, 16) and W2 is (16, 1), so Z2 is (300, 1) — one logit per sample.

Applies sigmoid to convert logits to probabilities. squeeze() removes the size-1 last dimension: (300, 1) → (300,). Each value is now P(class=1) for that sample.

Cross-entropy is the correct loss function for classification. Unlike MSE (which treats outputs as regression targets), cross-entropy measures the information-theoretic distance between the predicted probability distribution and the true label. It produces stronger gradients when predictions are confidently wrong.

A tiny constant to prevent log(0) = -infinity. Since A2 can be extremely close to 0 or 1, we add eps before taking the log. This is standard practice in every ML framework.

Binary Cross-Entropy (BCE) loss: L = -(1/N) ∑ [yᵢ · log(aᵢ) + (1-yᵢ) · log(1-aᵢ)]. For each sample, only one term is active: if y=1, the loss is -log(a) (penalizes low predictions); if y=0, the loss is -log(1-a) (penalizes high predictions). The minimum is 0 when every prediction perfectly matches its label.

Converts probabilities to binary predictions (>0.5 → class 1, ≤0.5 → class 0), then compares with true labels. The mean of this boolean array is the fraction of correct predictions.

Backpropagation computes the gradient of the loss with respect to every weight. The chain rule flows the error signal from the output back through each layer. For BCE + sigmoid, the output gradient has the elegant form: dZ2 = (A2 - y) / N.

The gradient of BCE loss with respect to the output logit Z2. The math works out beautifully: d(BCE)/d(Z2) = (sigmoid(Z2) - y) / N = (A2 - y) / N. The /300 averages over all samples. Reshape from (300,) to (300, 1) for matrix math.

Gradient for W2: the matrix product of hidden activations (transposed) and the output gradient. H.T is (16, 300) and dZ2 is (300, 1), giving dW2 of shape (16, 1) — matching W2’s shape.

Bias gradient is the sum of output gradients across all samples. Shape (1,) — matches b2.

Sends the error signal back through W2 to the hidden layer. dZ2 is (300, 1) and W2.T is (1, 16), giving dH of shape (300, 16). Each hidden neuron receives a portion of the output error, weighted by its connection strength W2.

Applies the ReLU derivative: gradient passes through active neurons (Z1 > 0) and is blocked by dead neurons (Z1 ≤ 0). The expression (Z1 > 0) creates a boolean mask that is 1 where ReLU was active and 0 where it was not.

Same pattern as dW2: gradient = input.T @ layer_gradient. X.T is (2, 300) and dZ1 is (300, 16), giving dW1 of shape (2, 16) — matching W1.

Sum of hidden gradients across all 300 samples. Shape (16,) matches b1.

Gradient descent: move weights in the direction that reduces loss. W_new = W_old - lr × gradient. We update the output layer first, but the order does not matter since we already computed all gradients.

Same update for the hidden layer. After this line, all 65 parameters have been updated. One epoch is complete.

Print metrics every 100 epochs and at the final epoch. This lets us track convergence without flooding the output.

Prints the epoch number, BCE loss, and classification accuracy. The :3d formats the epoch as a 3-digit integer, .4f gives 4 decimal places, .1% gives one decimal percentage.

Epoch   0: BCE=0.4991  acc=77.3%
Epoch 100: BCE=0.2186  acc=89.7%
Epoch 200: BCE=0.1079  acc=97.3%
Epoch 300: BCE=0.0633  acc=98.3%
Epoch 400: BCE=0.0463  acc=99.0%
Epoch 499: BCE=0.0377  acc=99.0%

PyTorch’s core library. Provides tensors (GPU-accelerated arrays), automatic differentiation (autograd), and the computational graph that enables loss.backward() to compute all gradients in one call.

The neural network module. Contains layer types (nn.Linear, nn.Conv2d), activations (nn.ReLU, nn.GELU), loss functions (nn.CrossEntropyLoss), containers (nn.Sequential), and the nn.Module base class that all networks must extend.

DataLoader handles batching, shuffling, and iterating over data. TensorDataset wraps tensors into a dataset object that DataLoader can consume. Together, they replace our manual for-loop over samples.

Every PyTorch neural network is a class that extends nn.Module. This is the fundamental pattern: define layers in __init__(), define the data flow in forward(). nn.Module gives us automatic parameter tracking, GPU support, save/load, and gradient computation for free.

Constructor that defines the network architecture. Takes flexible parameters so one class can build any MLP: MLP(2, [16], 1) for a 2→16→1 network, MLP(2, [8, 8], 1) for a 2→8→8→1 network.

Calls the parent class (nn.Module) constructor. This sets up PyTorch’s internal parameter tracking, hooks, and module registration. You MUST call this before defining any layers — omitting it causes subtle bugs.

We accumulate layer objects in a list, then pass them to nn.Sequential. This lets us build networks of any depth from the hidden_dims parameter.

The input dimension for the next layer. Starts at in_dim (2), then updates to each hidden width. For MLP(2, [16], 1): prev goes 2 → 16.

Iterates over hidden layer widths. For hidden_dims = [16]: one iteration creating Linear(2, 16) + ReLU. For [8, 8]: two iterations creating Linear(2, 8) + ReLU + Linear(8, 8) + ReLU.

Creates a fully-connected layer mapping prev inputs to h outputs. Internally stores weight matrix W of shape (h, prev) and bias vector b of shape (h,). The forward pass computes: output = input @ W.T + b.

Adds ReLU after each hidden layer. nn.ReLU() has no learnable parameters — it just applies max(0, x) element-wise. We add it after every hidden Linear but NOT after the output Linear.

The output of this layer (h neurons) becomes the input to the next layer. For hidden_dims = [8, 8]: after first iteration prev = 8, so the second Linear is Linear(8, 8).

Creates the final layer: prev inputs → out_dim outputs. For binary classification with out_dim = 1: this produces a single logit (raw score before sigmoid). No activation is added here — we use BCEWithLogitsLoss which applies sigmoid internally.

Wraps all layers into an nn.Sequential container. The * unpacks the list. When you call self.net(x), it passes x through each layer in order. Assigning to self.net (a nn.Module attribute) registers all sub-layers so model.parameters() finds them.

Calls our custom initialization method. PyTorch’s default is Kaiming uniform, which is close to He init. We explicitly apply Kaiming normal for consistency with our NumPy code.

Iterates through all layers and applies He (Kaiming) initialization to Linear layers. The leading underscore _ is a Python convention indicating this is an internal method not meant to be called externally.

self.modules() returns every module in the model tree: the MLP itself, the Sequential, each Linear, and each ReLU. We filter for nn.Linear to only initialize weight-bearing layers.

Only initialize nn.Linear layers (which have weights). nn.ReLU has no parameters, so initializing it would be meaningless.

He initialization: fills weights from N(0, √(2/fan_in)). The trailing _ means in-place modification. nonlinearity=‘relu’ tells PyTorch to use the factor of 2 (because ReLU zeroes half the activations).

Sets all biases to 0. Same as our b1 = np.zeros(16) in NumPy.

The forward method defines how data flows through the network. PyTorch calls this automatically when you do model(x). You NEVER call model.forward(x) directly — model(x) also triggers hooks and autograd tracking.

Passes x through the Sequential container: Linear → ReLU → Linear. The output is raw logits (no sigmoid). For a batch of 32: input (32, 2) → hidden (32, 16) → output (32, 1).

We split the 300-point dataset into 240 training and 60 validation samples. Training data is used to update weights; validation data is held out to detect overfitting — when the model memorizes training data but fails on new data.

Converts the first 240 NumPy rows to a PyTorch tensor. float32 is required because nn.Linear expects 32-bit floats (not 64-bit doubles or integers).

Training labels as float32 tensor. BCEWithLogitsLoss requires float targets, not integers. Shape: (240,).

The last 60 samples for validation. These are never used for weight updates — only for measuring how well the model generalizes to unseen data.

Validation labels. Shape: (60,). Used only in the evaluation phase.

Creates a mini-batch data pipeline. TensorDataset bundles (X_t, y_t) so dataset[i] = (X_t[i], y_t[i]). DataLoader splits the 240 samples into batches of 32, shuffled each epoch. This gives 8 batches per epoch (7 full batches of 32 + 1 batch of 16).

We encapsulate training in a function so we can easily compare different architectures. Each call creates a fresh optimizer, trains for 300 epochs, and returns final accuracies.

Takes a model, trains it, and returns train/validation accuracy. Separating training into a function is good practice — it avoids copy-pasting the training loop for each architecture.

Creates an Adam optimizer for this model. Adam maintains per-parameter momentum (m) and squared-gradient (v) buffers, then updates: w -= lr × m̂/(√v̂ + ε). This adaptive rate lets Adam converge faster than vanilla SGD on most problems.

Binary Cross-Entropy loss with built-in sigmoid. Unlike calling sigmoid() + BCE separately, BCEWithLogitsLoss uses the log-sum-exp trick for numerical stability. It expects raw logits (not probabilities) as input.

Each epoch iterates over all mini-batches in the DataLoader. With batch_size=32 and 240 samples, that is 8 optimizer steps per epoch, 2400 total steps over 300 epochs.

Switches the model to training mode. This enables dropout (if any) and uses running statistics for batch normalization. Our simple MLP has no dropout or batch norm, but it is good practice to always call model.train() before training.

The DataLoader yields one batch at a time: (xb, yb) where xb is shape (32, 2) and yb is shape (32,). The last batch may be smaller (16 samples if 240 is not divisible by 32).

Passes the batch through the model. model(xb) calls forward(), which returns shape (32, 1). squeeze(-1) removes the last dimension: (32, 1) → (32,) to match yb’s shape.

Computes BCE loss over the batch. The loss function applies sigmoid internally, then computes cross-entropy, then averages over the batch. Returns a scalar tensor with gradient tracking.

Resets all parameter gradients to zero. PyTorch accumulates gradients by default (adding new gradients to existing ones). Without this, gradients from batch 1 would carry into batch 2, corrupting the update.

Computes gradients for ALL model parameters via reverse-mode automatic differentiation. This single call replaces our entire NumPy backward pass (dZ2, dW2, db2, dH, dZ1, dW1, db1). PyTorch traces the computational graph from loss back to every weight.

Applies Adam’s update rule to every parameter: updates momentum m, updates squared gradient v, applies bias correction, then updates the weight. This is the equivalent of our W -= lr*dW but with adaptive per-parameter rates.

Switches to evaluation mode: disables dropout and uses running statistics for batch norm. For our simple MLP, the behavior is identical to train mode, but always switch modes as a habit.

Inside this block, PyTorch skips building the computational graph for gradient computation. This saves memory and speeds up inference. We don’t need gradients during evaluation — only during training.

Computes training accuracy in one line. The logit > 0 check is equivalent to sigmoid > 0.5 (sigmoid(0) = 0.5). We compare predictions with labels, convert booleans to float (True=1, False=0), and take the mean.

Same accuracy computation on the 60 validation samples. If val_acc << train_acc, the model is overfitting (memorizing training data). If both are high, the model generalizes well.

Returns accuracy as plain Python floats. .item() converts a single-element tensor to a float. This is needed for printing — f-strings work better with Python floats than torch tensors.

Now we systematically test five architectures: three single-hidden-layer networks (varying width) and two two-hidden-layer networks (varying width). This reveals how architecture affects generalization.

Maps architecture names to hidden layer dimension lists. The name uses → arrows to show data flow. Each list is passed to MLP() as hidden_dims.

Sets PyTorch’s random seed before the comparison loop. This ensures each model gets a consistent random initialization, making the comparison fair.

Iterates over all five architectures. For each: builds the model, counts parameters, trains for 300 epochs, and evaluates.

Creates a fresh MLP instance with the specified architecture. Each call to MLP() runs __init__ which builds the layers and applies He initialization.

Counts total learnable parameters using a generator expression. model.parameters() yields each weight/bias tensor, p.numel() counts elements in each.

Calls our training function. Returns train and validation accuracy after 300 epochs of Adam optimization.

Prints the summary line. :14s pads the name to 14 chars, :4d formats parameter count, :.1% formats accuracy as a percentage with one decimal.

2→4→1        params=  17  train=100.0%  val=96.7%
2→16→1       params=  65  train=100.0%  val=98.3%
2→32→1       params= 129  train=100.0%  val=98.3%
2→8→8→1      params= 105  train=100.0%  val=98.3%
2→16→16→1    params= 337  train=100.0%  val=98.3%