← Foundation lessonsFoundations · lesson 2

Neural network: one input

In PyTorch basics the model was one straight line. Much of the world isn’t straight: the arc of a thrown ball, or a speed that flips at a bounce.

This lesson builds the smallest model that can bend, a small neural network, one line of PyTorch at a time.

Overview

Here is the line this lesson explains. By the end, every part of it will be familiar:

The goal
network = nn.Sequential(nn.Linear(1, 4), nn.ReLU(), nn.Linear(4, 1))
A linear layernn.Linear with more outputs.An activationnn.ReLU: the bend.Neurons and layersnn.Sequential.Fit a curveThe same training loop.

The training loop does not change: guess, loss, zero_grad(), backward(), step(). Only the model gets more interesting.

Linear layer

nn.Linear(in, out) takes in numbers and gives back out numbers.

Linear layer
what nn.Linear is called: several weighted sums of the same inputs, side by side.

The two numbers are not the weight and bias counts. Each output gets:

  • its own weight for every input: in × out weights;
  • one bias: out biases.

In nn.Linear(1, 1) both counts happen to be 1. Bigger layers show the difference:

LayerInOutWeightsBiasesWhat it computes
nn.Linear(1, 1)1111one output: \(wx + b\)
nn.Linear(1, 4)1444four outputs from the same \(x\): \(w_j x + b_j\)
nn.Linear(4, 1)4141one output from four inputs: \(w_1 x_1 + w_2 x_2 + w_3 x_3 + w_4 x_4 + b\)
nn.Linear(4, 16)416641616 outputs, each from all four inputs

PyTorch stores the weights as a table: one row per output, one column per input.

examples/pytorch_small_network.py · a bigger linear layerOpen source ↗
layer = nn.Linear(4, 16)
print(layer.weight.shape, layer.bias.shape)   # torch.Size([16, 4]) torch.Size([16])
Full file · pytorch_small_network.py · 45 lines
"""A small network: fit a curve that no straight line can follow.

Run: python3 examples/pytorch_small_network.py
Uses the training loop from pytorch_first_update.py unchanged; only the model (the network) is new.
"""

import torch
from torch import nn

x = torch.linspace(-1, 1, 41).reshape(-1, 1)   # 41 inputs, one per row: shape (41, 1)
y = 1 - (x - 0.2) ** 2                         # the curve we want the network to learn


def make_network(hidden):
    if hidden == 0:
        return nn.Linear(1, 1)                 # no hidden layer: the straight line from PyTorch basics
    return nn.Sequential(nn.Linear(1, hidden), nn.ReLU(), nn.Linear(hidden, 1))


def train(hidden, seed, steps=3000):
    torch.manual_seed(seed)                    # the random starting weights
    network = make_network(hidden)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.SGD(network.parameters(), lr=0.1)
    for _ in range(steps):
        loss = loss_fn(network(x), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return network, loss_fn(network(x), y).item()


if __name__ == "__main__":
    layer = nn.Linear(4, 16)
    print("nn.Linear(4, 16) weight", tuple(layer.weight.shape), "bias", tuple(layer.bias.shape))
    print("ReLU of [-2, -0.5, 0, 1.5]:", torch.relu(torch.tensor([-2.0, -0.5, 0.0, 1.5])).tolist())

    network = nn.Sequential(nn.Linear(1, 4), nn.ReLU(), nn.Linear(4, 1))
    first, relu, second = network               # nn.Sequential holds the three pieces in order
    print("network(x) == second(relu(first(x))):", torch.equal(network(x), second(relu(first(x)))))
    print("first.weight", tuple(first.weight.shape), "first.bias", tuple(first.bias.shape))   # one row per hidden neuron

    for hidden in (0, 1, 2, 4, 16):
        losses = [train(hidden, seed)[1] for seed in (0, 1, 2)]
        print(f"{hidden:2d} hidden neurons, starts 1-3: loss " + "  ".join(f"{v:.4f}" for v in losses))

Each output \(j\) is a weighted sum of the inputs, plus a bias:

\[ z_j = w_{j1} x_1 + w_{j2} x_2 + w_{j3} x_3 + w_{j4} x_4 + b_j \]

The weights of nn.Linear(4, 16): one row per output, plus one bias per row.

Row 1 holds \(w_{11}, w_{12}, w_{13}, w_{14}\) and \(b_1\), and turns the four inputs into output 1. In plain words: 16 straight-line rules side by side, computed at once.

That is \(16 \times 4 = 64\) weights plus 16 biases: 80 adjustable numbers.

Weighted sum
each input times its weight, added up; the bias is added on top.

From here on we keep a single input, \(x\), so everything fits on one graph. (Part 2 brings back more inputs.)

With one input, nn.Linear(1, 4) gives four outputs. Each is a straight line \(z_j = w_j x + b_j\) with its own weight and bias.

Activation

Adding straight lines together only ever makes another straight line. To get a curve, we need something that bends.

Activation function
a fixed function applied to each output of a linear layer.
ReLU
rectified linear unit, the most common activation: keep a positive number, turn a negative one into zero.

\[ \mathrm{ReLU}(z) = \max(0,\, z) \]

ReLU on four numbers
torch.relu(torch.tensor([-2.0, -0.5, 0.0, 1.5]))   # tensor([0.0, 0.0, 0.0, 1.5])

Apply ReLU to a straight line \(wx + b\), and the part below zero is flattened.

What’s left is a ramp: flat on one side, sloped on the other. Its bend sits where the line crosses zero, at \(x = -b/w\).

Move the sliders: the code, its output and the five dots update together.

One ramp
Live · ReLU of a line, at five inputs

nn.ReLU() has no adjustable numbers: it is a fixed function. The bend comes from the line before it, and training moves that line.

Neurons and layers

Neuron
a weighted sum plus a bias, followed by an activation. Also called a unit.

With inputs \(x_1, \dots, x_n\), weights \(w_1, \dots, w_n\) and bias \(b\), a ReLU neuron computes:

\[ a = \mathrm{ReLU}(w_1 x_1 + w_2 x_2 + \dots + w_n x_n + b) \]

With one input, that is \(a = \mathrm{ReLU}(wx + b)\): exactly the ramp you just dragged.

Layer
a group of neurons that work side by side on the same inputs.

Our network has two layers:

  • a hidden layer of four ReLU neurons;
  • an output layer of one neuron, which adds their ramps into the answer.

The input and the output are not layers. They are just numbers going in and coming out.

The network: one input, four hidden ReLU neurons, one output neuron.

How to read it:

  • Each line is one weight.
  • Each hidden circle is a neuron: it adds its weighted input and its bias, then applies ReLU.
  • The Σ circle is the output neuron. It adds the four ramps with weights \(v_1, \dots, v_4\) and bias \(c\), with no ReLU.
  • Biases are not drawn.

Each neuron owns the lines coming into it, plus its bias.

  • A hidden neuron owns 1 weight and 1 bias: 4 neurons, 8 numbers.
  • The output neuron owns its 4 incoming weights and 1 bias: 5 numbers.

Lines leaving a neuron belong to the next neuron.

The same thing as a table, with the PyTorch pieces beside the diagram names:

NameWhat it isIn PyTorchParameters
Inputthe numbers you feed in, \(x\)the tensor you pass to network(x); no module0
Hidden layer4 neurons, each a weighted sum + ReLUnn.Linear(1, 4) + nn.ReLU(): two modules, one layer8
Output layer1 linear neuron: a weighted sum, no activationnn.Linear(4, 1)5
Outputthe number it produces, \(y\)the tensor network(x) returns0

Two pairs of words to keep apart:

  • The output layer is the last computing part. The output is the number it produces.
  • In PyTorch, “layer” is used loosely for any module. nn.Sequential(nn.Linear(1, 4), nn.ReLU(), nn.Linear(4, 1)) has three modules, but two layers in diagram terms: one hidden, one output.

Why does the output layer have no ReLU? The two layers have different jobs.

  • The hidden layer makes the bends, and that needs the activation.
  • The output layer only adds the ramps into the final answer, and that answer must be able to take any value.

A ReLU at the output would make negative predictions impossible. The bouncing-ball workshop has to predict a speed change of \(-0.098\) on every step in the air.

It would also stall training: whenever the sum came out negative, the output would sit at 0 with zero gradient.

Regression
predicting a number. Its standard output layer is linear, with no activation.

Models that pick between categories use softmax at the output instead; that comes in Neural network: several inputs.

nn.Sequential runs the pieces in order, feeding each result into the next. Here it is, with the same calculation written out by hand:

examples/pytorch_small_network.py · a network, two waysOpen source ↗
network = nn.Sequential(nn.Linear(1, 4), nn.ReLU(), nn.Linear(4, 1))
first, relu, second = network               # nn.Sequential holds the three pieces in order
print(torch.equal(network(x), second(relu(first(x)))))   # True
print(first.weight.shape, first.bias.shape)  # torch.Size([4, 1]) torch.Size([4]): one row per neuron

Full file: pytorch_small_network.py, in its first card above ↑

first = nn.Linear(1, 4)hidden layer, part 1: \(z_j = w_j x + b_j\) for the 4 neurons · 8 numbers
relu = nn.ReLU()hidden layer, part 2: \(a_j = \mathrm{ReLU}(z_j)\), four ramps · 0 numbers
second = nn.Linear(4, 1)output layer: \(y = v_1 a_1 + v_2 a_2 + v_3 a_3 + v_4 a_4 + c\) · 5 numbers

So the network’s answer is a bias plus four weighted ramps.

Each ramp is flat up to its bend and straight after it. Added up, they make a line of straight pieces that changes direction at every bend: more neurons, more bends.

PyTorch has no Neuron class. It computes a whole layer at once, because one table calculation is far faster than a loop over neurons.

A neuron still exists in the code, as a slice: hidden neuron \(j\) is row \(j\) of first.weight, plus first.bias[j], plus the ReLU on output \(j\).

With \(h\) hidden neurons, one input and one output, the network has \(2h + (h + 1) = 3h + 1\) adjustable numbers.

With a batch of 41 inputs, the shapes go (41, 1) → (41, 4) → (41, 4) → (41, 1).

Curve fitting

Now the test. The target is an arc, \(y = 1 - (x - 0.2)^2\), at 41 points from −1 to 1. No straight line can follow it.

The training code is the PyTorch basics loop, word for word. Only make_network is new:

examples/pytorch_small_network.py · trainingOpen source ↗
torch.manual_seed(seed)                    # the random starting weights
network = make_network(hidden)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(network.parameters(), lr=0.1)
for _ in range(3000):
    loss = loss_fn(network(x), y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Full file: pytorch_small_network.py, in its first card above ↑

We trained networks with 0, 1, 2, 4 and 16 hidden neurons, each from three random starts. Zero hidden neurons is just nn.Linear(1, 1), the straight line.

Pick one. The plot shows the network’s answer, its ramps and its bends. The diagram beside it shows the same network:

  • a thicker line is a bigger weight;
  • blue is a positive weight, magenta a negative one;
  • a grey neuron is switched off.
Compare networks
Hidden neurons
Random start
Adjustable numbers—
Loss after 3000 steps—
Bends inside the data—

What the runs show:

  • 0 neurons, the straight line: loss 0.098, the best a line can do.
  • 1 neuron adds one bend. From start 1, that alone cuts the loss to 0.022.
  • 2 and 4 neurons do better still, if the start is good.
  • 16 neurons follow the arc closely from every start (0.0001 to 0.0012).

Why do small networks sometimes get stuck? A neuron whose weighted sum is below zero for every input outputs zero everywhere. Nudging its weights then changes nothing: its gradient is zero, and it never recovers.

Dead neuron
a ReLU neuron that outputs zero for every input, so training cannot revive it. The “dying ReLU” problem.

The next lesson, Gradients and backpropagation, shows it from the inside.

A bend outside −1…1 is no help either. Inside the data, that ramp is just a straight line.

So what does the number of neurons mean? It is room, not a list of parts.

Nobody gives a neuron its job; training does. Here, a neuron's job is its ramp.

Pick 16 and count: in our three runs, only 5 to 7 neurons bend inside the data. The rest are switched off or bend outside −1…1.

So 16 is not 16 jobs. It is enough spare room that every random start still finds the few bends the arc needs. With 4, one start found none.

That is how to choose a size: try a few, and take the smallest that works from every start. The bouncing-ball workshop chooses its size the same way.

Recap

  • nn.Linear(4, 16) is a linear layer: 16 weighted sums of the 4 inputs, each with its own bias, stored as a 16 × 4 weight table plus 16 biases.
  • nn.ReLU() is an activation function, \(\max(0, z)\): it turns each straight line into a ramp with one bend. It has nothing to train.
  • A neuron is one weighted sum plus bias, followed by the activation: \(\mathrm{ReLU}(w_1x_1 + \dots + w_nx_n + b)\). In PyTorch it is one row of a layer, not an object of its own.
  • nn.Sequential(...) runs its layers in order: network(x) is second(relu(first(x))).
  • nn.Linear(h, 1) at the end is a linear output neuron: it weights the \(h\) ramps and adds them up, so the answer can bend once per hidden neuron.
  • The training loop is unchanged: loss_fn, zero_grad(), backward(), step() adjust every weight and bias in every layer at once.

Read the goal line again: nn.Sequential(nn.Linear(1, 4), nn.ReLU(), nn.Linear(4, 1)) is a hidden layer of four ReLU neurons and one linear output neuron.

Four ramps, added up into one answer that can follow a curve.

Key concepts

Linear layer
Many weighted sums of the same inputs, each with its own bias. nn.Linear(in, out) has an out × in table of weights and one bias per output.
Activation function
A fixed function applied to each output of a linear layer. ReLU, \(\max(0, z)\), is the most common; it makes a ramp with one bend.
Neuron (unit)
A weighted sum plus a bias, followed by an activation: \(\mathrm{ReLU}(w_1x_1 + \dots + w_nx_n + b)\). An output neuron with no activation is a linear neuron.
Layer · hidden layer
A group of neurons that read the same inputs. Hidden layers sit between the input and the output.
Neural network
Layers of neurons connected in order. Ours: 1 input, a hidden layer of ReLU neurons, 1 linear output neuron.
nn.Sequential
A container that runs its modules in order, feeding each result into the next.
Weighted sum
Each input times its weight, added up; a neuron adds its bias on top.
Regression
Predicting a number. The standard output layer for it is linear.
Input layer
What some books call the input column. It computes nothing and has no parameters, so this lesson does not count it as a layer.
Dead neuron
A ReLU neuron whose weighted sum is negative for every input. Its gradient is zero, so training cannot revive it.

Every curve on this page is a network trained by the example file and saved by tools/build_network_data.py; the page only draws them.