Layers and ReLU

In diagrams a neural network looks like a nervous system. In code it is a list. You call one object after another, hand each the result of the previous one and get a prediction at the end. That is the whole of the magic.

Two building blocks are enough for a network that can do more than a straight line: a linear layer, which multiplies by a weight matrix and adds a bias, and an activation, which does something non-linear. Without the second one a stack of linear layers would mathematically be a single one - three matrices in a row are just another matrix.

The goal

A network is a chain of simple parts. Every part can do exactly two things, and you have already built both without noticing:

Linear.forward is x @ W + b. Those are the two operators from chapter 01. You do not have to program the backward direction anywhere - it falls out of the autograd from chapter 02, because @ and + bring it along. That is exactly why autograd came so early on the roadmap.

ReLU is the simplest usable non-linearity: everything below zero becomes zero, the rest stays. It is the one place in this chapter where your tensor has to learn a new operation - and its derivative is a switch: one where the input was positive, zero otherwise.

Zero magic in this step

NumPy is allowed, torch, jax, tinygrad and autograd are forbidden. The temptation is especially tangible in this chapter, because nn.Linear and nn.ReLU are lying around ready-made in PyTorch. But those two classes together are not thirty lines - and anyone who has typed them once knows forever what model.parameters() actually returns: a flat list of the tensors the optimizer is later allowed to turn the screws on.

The exercise

Two files. First a new method in meintorch/tensor.py:

    def relu(self):
        out = Tensor(np.maximum(self.data, 0.0), (self,), "relu")

        def _backward():
            # TODO(du): the gradient only flows on where the input was
            # positive. Everywhere else it is blocked.
            raise NotImplementedError

        out._backward = _backward
        return out

Then meintorch/nn.py:

import numpy as np

from .tensor import Tensor


class Module:
    """Shared base: be callable and hand out your parameters."""

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

    def forward(self, x):
        raise NotImplementedError

    def parameters(self):
        return []


class Linear(Module):
    def __init__(self, in_features, out_features, seed=None):
        # Uniform around zero, scaled by the input width. For free -
        # initialisation is a topic of its own and not this chapter.
        rng = np.random.default_rng(seed)
        limit = 1.0 / np.sqrt(in_features)
        self.W = Tensor(rng.uniform(-limit, limit, (in_features, out_features)))
        self.b = Tensor(np.zeros(out_features))

    def forward(self, x):
        # TODO(du): one line - matrix product with the weights, bias on top.
        raise NotImplementedError

    def parameters(self):
        return [self.W, self.b]


class ReLU(Module):
    def forward(self, x):
        # TODO(du): one line as well - the tensor can do this itself now.
        raise NotImplementedError


class Sequential(Module):
    def __init__(self, *layers):
        self.layers = list(layers)

    def forward(self, x):
        # TODO(du): every layer in turn, the result travels onward.
        raise NotImplementedError

    def parameters(self):
        return [p for layer in self.layers for p in layer.parameters()]

Hint 1 - the direction

All three forward methods are one to three lines. In this chapter you write not a single derivative apart from ReLU's - everything else is handled by the graph you already have. If you notice yourself applying the chain rule by hand, you are on the wrong track.

For Sequential.forward: it is a loop that keeps overwriting one variable.

Hint 2 - more concrete

• ReLU backward step: (self.data > 0) is a mask of True and False. Multiplied by out.grad it becomes exactly what you need. Do not forget to accumulate.
• Linear.forward: x @ self.W + self.b. Nothing else.
• ReLU.forward: x.relu().
• Sequential.forward: for layer in self.layers: x = layer(x) and return x at the end.

    # in Tensor.relu
        def _backward():
            self.grad += (self.data > 0) * out.grad


    # in nn.py
    class Linear(Module):
        def forward(self, x):
            return x @ self.W + self.b


    class ReLU(Module):
        def forward(self, x):
            return x.relu()


    class Sequential(Module):
        def forward(self, x):
            for layer in self.layers:
                x = layer(x)
            return x

In goal mode

The autonomous run builds both files and proves correctness numerically instead of relying on the shape of the output:

Extend `meintorch/tensor.py` with a method `relu()` and its matching backward
step, and create `meintorch/nn.py` with the classes `Module`, `Linear`, `ReLU`
and `Sequential`. `Linear` initialises W uniformly with ±1/sqrt(in_features), b
with zeros, and hands both out via `parameters()`. `Sequential` passes the input
through all layers and collects their parameters.
Mandatory: NumPy allowed, torch/jax/tinygrad/autograd forbidden. No derivative
may be written by hand outside `tensor.py`.
Write `tests/test_03_schichten.py`: forward and backward have to run through a
network of two linear layers with a ReLU in between, and every parameter has to
carry a gradient in its own shape afterwards. Check at least one weight against
a central numerical difference.
Run `python -m pytest` until everything is green.

The milestone test

tests/test_03_schichten.py. The last test is the interesting one: it compares the gradient your autograd computes with the numerical derivative at the same spot. If those two agree, the entire chain is correct - from the matrix product through ReLU and back into the weights.

import numpy as np

from meintorch.nn import Linear, ReLU, Sequential
from meintorch.tensor import Tensor


def network_and_input():
    net = Sequential(Linear(3, 4, seed=0), ReLU(), Linear(4, 1, seed=1))
    x = Tensor(np.random.default_rng(7).normal(size=(5, 3)))
    return net, x


def test_forward_returns_the_right_shape():
    net, x = network_and_input()
    assert net(x).shape == (5, 1)


def test_relu_cuts_off_negative_values():
    out = ReLU()(Tensor([[-2.0, 0.0, 3.0]]))
    assert np.allclose(out.data, [[0.0, 0.0, 3.0]])


def test_backward_reaches_every_parameter():
    net, x = network_and_input()
    net(x).backward()
    for p in net.parameters():
        assert p.grad.shape == p.data.shape
    assert np.any(net.layers[2].W.grad != 0.0)


def test_gradient_matches_the_numerical_derivative():
    net, x = network_and_input()
    net(x).backward()
    W = net.layers[2].W
    analytic = W.grad[0, 0]

    eps = 1e-6
    W.data[0, 0] += eps
    plus = net(x).data.sum()
    W.data[0, 0] -= 2 * eps
    minus = net(x).data.sum()
    W.data[0, 0] += eps

    assert abs(analytic - (plus - minus) / (2 * eps)) < 1e-4