Skip to content

PigeonColumba compressaTheme

Lesson 3 of 10

Train a two-layer network from scratch

Put lessons 1 and 2 together: a network x -> relu(x W1 + b1) -> logits = h W2 + b2 trained with plain mini-batch gradient descent on a two-spirals dataset (provided). No libraries beyond numpy.

The loop is the same one lab/train.py:fit runs, minus Adam and extras:

Shuffle, take a batch.
Forward: hidden activations, logits.
dz = (softmax(logits) - onehot(y)) / N.
Backward through the output layer, then the ReLU layer.
param -= lr * grad for all four parameters.

Implement train(X, y, hidden, epochs, lr, seed) returning a predict(X) function. The test requires 90% training accuracy on the spirals, which a linear model cannot get (about 50%), so your hidden layer has to work.

Your code

import numpy as np


def spirals(n=400, seed=0):
    rng = np.random.default_rng(seed)
    t = np.sqrt(rng.random(n)) * 3 * np.pi
    y = rng.integers(0, 2, n)
    r = t * (2 * y - 1)
    X = np.stack([r * np.cos(t), r * np.sin(t)], 1) / 10 + rng.normal(0, 0.05, (n, 2))
    return X, y


def train(X, y, hidden=64, epochs=300, lr=0.5, seed=0):
    """Train a 1-hidden-layer ReLU network with mini-batch SGD.
    Return a function predict(X) -> class indices."""
    rng = np.random.default_rng(seed)
    C = int(y.max()) + 1
    W1 = rng.standard_normal((X.shape[1], hidden)) * np.sqrt(2 / X.shape[1])
    b1 = np.zeros(hidden)
    W2 = rng.standard_normal((hidden, C)) * np.sqrt(2 / hidden)
    b2 = np.zeros(C)
    raise NotImplementedError

Checking your code and saving a draft need JavaScript, which is turned off in this browser. Turn it on and reload the page to check your code. You can still edit it here, and the hints and the reference solution are below.

Pigeon's hints (2)

Hint 1 of 2

Reuse your softmax and dense-layer code from lessons 1 and 2.

Hint 2 of 2

If accuracy is stuck near 0.5, check the ReLU mask uses z1 (pre-activation), and that you update W1/b1, not only W2/b2.

The reference solution

import numpy as np


def spirals(n=400, seed=0):
    rng = np.random.default_rng(seed)
    t = np.sqrt(rng.random(n)) * 3 * np.pi
    y = rng.integers(0, 2, n)
    r = t * (2 * y - 1)
    X = np.stack([r * np.cos(t), r * np.sin(t)], 1) / 10 + rng.normal(0, 0.05, (n, 2))
    return X, y


def train(X, y, hidden=64, epochs=300, lr=0.5, seed=0):
    rng = np.random.default_rng(seed)
    C = int(y.max()) + 1
    W1 = rng.standard_normal((X.shape[1], hidden)) * np.sqrt(2 / X.shape[1])
    b1 = np.zeros(hidden)
    W2 = rng.standard_normal((hidden, C)) * np.sqrt(2 / hidden)
    b2 = np.zeros(C)
    for _ in range(epochs):
        idx = rng.permutation(len(X))
        for s in range(0, len(X), 32):
            j = idx[s:s + 32]
            x, t = X[j], y[j]
            z1 = x @ W1 + b1
            h = np.maximum(z1, 0)
            z2 = h @ W2 + b2
            p = np.exp(z2 - z2.max(1, keepdims=True))
            p /= p.sum(1, keepdims=True)
            p[np.arange(len(t)), t] -= 1
            dz2 = p / len(t)
            dW2, db2 = h.T @ dz2, dz2.sum(0)
            dz1 = (dz2 @ W2.T) * (z1 > 0)
            dW1, db1 = x.T @ dz1, dz1.sum(0)
            W1 -= lr * dW1; b1 -= lr * db1; W2 -= lr * dW2; b2 -= lr * db2

    def predict(Xn):
        return (np.maximum(Xn @ W1 + b1, 0) @ W2 + b2).argmax(1)
    return predict

Where this is used

Regularizers & training tricksRuns hereLabel smoothing, mixup, weight averaging (EMA), longer training: what each does for a capacity-starved model.