A Single Neuron (from Scratch)

Code

Neuron and softmax — fill in two holes

In the last step you saw the neuron and softmax at work. Now you write them as Python — and a piece is missing from each.

Pseudocode

neuron(x, w, b):
    return  x @ w + b

softmax(z):                        # z is [n, 10], one row = one picture's scores
    for every row:
        e ← exp(score − the largest score in that row)
        return  e / (the sum of e in that row)

Nothing complicated, but there is one line that looks unnecessary — and is not.

"Subtract the largest." Mathematically softmax does not change: subtracting the same number from all scores leaves the ratios of the probabilities as they were. But on a computer exp(1000) is infinity, and infinity / infinity is nan. After subtracting the largest, the biggest score becomes 0, and exp(0) = 1 — it never overflows.

Your job

Create neuron.py and test_neuron.py (see the right-hand panel) in the same folder as tensors_demo.py.

neuron.py is complete and runs, but is wrong: two places are marked HOLE 1 and HOLE 2.

HOLE 1 — neuron. In place of the weighted sum there are "zeros of the right shape": the code runs, but every score is 0. Write the real formula.
HOLE 2 — softmax. The division by the sum is already there. What is missing is one line of stability: subtract each row's largest value before torch.exp.

Three hints for HOLE 2:

z.max(dim=1) returns a pair (values and indices); the values are .values;
"for every row" means dim=1;
the result must be [n, 1], not [n] — otherwise it subtracts along the wrong axis. Which max parameter keeps that, you remember from the lesson 2 drills (task 4).

sigmoid is already written — it is only there so you can see what one activation looks like.

Run the test before filling anything

$ python -m unittest
FFF.FF..
FAIL: test_bias_is_added (...)
AssertionError: 0.0 != 0.7 within 5 places (0.7 difference)
FAIL: test_many_neurons_many_images (...)
AssertionError: Lists differ: [[0.0, 0.0], [0.0, 0.0]] != [[1.5, 9.5], [5.5, 49.5]]
FAIL: test_one_neuron_one_image (...)
AssertionError: 0.0 != 4.75 within 5 places (4.75 difference)
FAIL: test_big_scores_do_not_overflow (...)
AssertionError: True is not false : got [[nan, 0.0, 0.0]] - overflow?
FAIL: test_every_row_is_stabilised_on_its_own (...)
AssertionError: True is not false : got [[nan, 0.0, 0.0], [nan, nan, nan]]
----------------------------------------------------------------------
Ran 8 tests in 0.003s

FAILED (failures=5)

(Only the important lines are kept.)

Notice what passes: three dots — sigmoid, "rows add up to 1" and "known softmax values". The unstable softmax works correctly on ordinary numbers and breaks only on big ones. A person would not notice this on their own — the test does: got [[nan, 0.0, 0.0]].

Fill both holes and run again — you should see OK.

When it does not pass

Each of these was tried here:

0.0 != 0.7 (the bias test), while the other two neuron tests are off by exactly the bias — you wrote x @ w and forgot + b. The formula is almost right: 4.5 != 4.75 and [[1.0, 10.0], [5.0, 50.0]] instead of [[1.5, 9.5], [5.5, 49.5]].
a Tensor with 3 elements cannot be converted to Scalar — you wrote x * w + b. * is element by element: you got three numbers, not one sum.
got [[1.0, 0.0, 0.0], [nan, nan, nan]] with z - z.max() — you subtracted one largest value for the whole table, not one per row. The first row survived; in the second exp(−2000) is 0, and 0 / 0 is nan.
The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1 — z.max(dim=1).values without keepdim=True: you got [n] where [n, 1] is needed, and PyTorch tries to match the wrong axis.

neuron.py CREATE

Create this file in ./

"""neuron.py — a neuron built by hand from raw tensor math. No nn.Module."""

import torch


def neuron(x: torch.Tensor, w: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """One neuron per column of w: a weighted sum of the inputs plus a bias."""
    # HOLE 1: the weighted sum of the inputs, plus the bias.
    # (This placeholder returns zeros of the right shape, so everything runs.)
    return torch.zeros(x.shape[:-1] + w.shape[1:])


def sigmoid(z: torch.Tensor) -> torch.Tensor:
    """Squash any number into (0, 1): 'how sure is this neuron?'"""
    return 1 / (1 + torch.exp(-z))


def softmax(z: torch.Tensor) -> torch.Tensor:
    """Turn a row of scores into probabilities that sum to 1."""
    # HOLE 2: subtract each row's largest score BEFORE exp(), so that big
    # scores cannot overflow. (Dividing by the sum below stays as it is.)
    e = torch.exp(z)
    return e / e.sum(dim=1, keepdim=True)

test_neuron.py CREATE

Create this file in ./

"""test_neuron.py — the verdict for neuron.py.

Run it from this folder:   python -m unittest
"""

import unittest

import torch

import neuron as nr


class TestNeuron(unittest.TestCase):
    def test_one_neuron_one_image(self):
        # 1*0.5 + 2*(-1) + 3*2 + 0.25 = 4.75
        x = torch.tensor([1.0, 2.0, 3.0])
        w = torch.tensor([0.5, -1.0, 2.0])
        b = torch.tensor(0.25)
        self.assertAlmostEqual(nr.neuron(x, w, b).item(), 4.75, places=5)

    def test_bias_is_added(self):
        x = torch.zeros(3)  # all-zero input: only the bias is left
        w = torch.ones(3)
        b = torch.tensor(0.7)
        self.assertAlmostEqual(nr.neuron(x, w, b).item(), 0.7, places=5)

    def test_many_neurons_many_images(self):
        # two images, three inputs each; two neurons side by side = [3, 2]
        x = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 1.0]])
        W = torch.tensor([[1.0, 10.0], [2.0, 20.0], [3.0, 30.0]])
        b = torch.tensor([0.5, -0.5])
        out = nr.neuron(x, W, b)
        self.assertEqual(tuple(out.shape), (2, 2))
        self.assertEqual(out.tolist(), [[1.5, 9.5], [5.5, 49.5]])


class TestSigmoid(unittest.TestCase):
    def test_sigmoid_squashes_into_zero_one(self):
        z = torch.tensor([-100.0, 0.0, 100.0])
        out = nr.sigmoid(z)
        self.assertAlmostEqual(out[1].item(), 0.5, places=6)
        self.assertTrue(bool((out >= 0).all() and (out <= 1).all()))


class TestSoftmax(unittest.TestCase):
    def test_rows_sum_to_one(self):
        z = torch.tensor([[1.0, 2.0, 3.0], [0.0, 0.0, 0.0]])
        sums = nr.softmax(z).sum(dim=1)
        self.assertTrue(torch.allclose(sums, torch.ones(2)))

    def test_known_values(self):
        out = nr.softmax(torch.tensor([[1.0, 2.0, 3.0]]))
        want = torch.tensor([[0.0900, 0.2447, 0.6652]])
        self.assertTrue(torch.allclose(out, want, atol=1e-4), f"got {out.tolist()}")

    def test_big_scores_do_not_overflow(self):
        out = nr.softmax(torch.tensor([[1000.0, 1.0, 2.0]]))
        self.assertFalse(bool(torch.isnan(out).any()), f"got {out.tolist()} - overflow?")
        self.assertAlmostEqual(out[0, 0].item(), 1.0, places=5)

    def test_every_row_is_stabilised_on_its_own(self):
        # row 1 is huge, row 2 is hugely negative: one shared maximum fails row 2
        z = torch.tensor([[1000.0, 1.0, 2.0], [-1000.0, -1001.0, -1002.0]])
        out = nr.softmax(z)
        self.assertFalse(bool(torch.isnan(out).any()), f"got {out.tolist()}")
        want = torch.tensor([0.6652, 0.2447, 0.0900])
        self.assertTrue(torch.allclose(out[1], want, atol=1e-4), f"row 2 is {out[1].tolist()}")


if __name__ == "__main__":
    unittest.main()