The loss works. Now — the whole loop that you saw in words in step 2. This time you write it without holes: only the skeleton and the contract.

Pseudocode

train(x, y, steps, lr, seed):
    set the random seed
    W ← small random weights [784, 10], marked "track the gradient"
    b ← zeros [10],                     marked "track the gradient"
    history ← an empty list

    repeat steps + 1 times:
        probs ← forward(x, W, b)
        loss  ← cross_entropy(probs, y)
        save (loss, accuracy) to history
        if this is the last time: stop         # here we only measure the result

        loss.backward()                        # fills W.grad and b.grad
        W ← W − lr · W.grad                    # the update, NOT tracked
        b ← b − lr · b.grad
        W.grad ← 0, b.grad ← 0                 # clean up

    return W, b, history

Notice that the loop runs steps + 1 times: the first measurement is before any update, and the last one after the final update. That is why history has steps + 1 entries.

What is new here

Three things you have to write yourself in Python, and each has its name in PyTorch:

• "Mark it to track the gradient" — requires_grad_() on a tensor (or requires_grad=True on creation, as in lesson 2). Without it backward() has nothing to fill.
• "The update is not tracked" — the update is not part of the computation we learn the gradient of. Write it inside a with torch.no_grad(): block. Without it PyTorch will try to track the update operation too, and that is not allowed on a "leaf" tensor.
• "Clean up" — W.grad.zero_(). Again lesson 2's fifth breakage: every backward() adds.

Write the update "in place": W -= lr * W.grad, not W = W - lr * W.grad. The difference is in the next step, but the test will catch it.

Your job

Create train_scratch.py and test_train_scratch.py (right panel).

train_scratch.py has one function, train, with a docstring and raise NotImplementedError. The whole contract is in the docstring; read it, not just the signature. Above — the imports from neuron.py.

test_train_scratch.py runs the training loop for 100 steps on 10,000 pictures and checks:

• the history has 101 entries;
• the starting loss ≈ ln 10, accuracy still low;
• every loss is a real number (not nan);
• the loss falls and ends below 1.0;
• after one step the accuracy has already risen by at least 5 points;
• after 100 steps accuracy is above 70 %;
• the returned weights have the right shape, and the bias has moved.

The first run (before you write anything) says only one thing:

$ python -m unittest
...................E
ERROR: setUpClass (test_train_scratch.TestTrain)
NotImplementedError
----------------------------------------------------------------------
Ran 19 tests in 0.089s

FAILED (errors=1)

19 dots are your neuron.py tests, passing. Now write train and run again. All the tests together should take a few seconds.

When it does not pass

All of these failures are real; some of them you will see again in the next step.

• RuntimeError: a leaf Variable that requires grad is being used in an in-place operation. — you wrote W -= ... outside a with torch.no_grad(): block.
• AttributeError: 'NoneType' object has no attribute 'zero_' — W = W - lr * W.grad instead of W -= lr * W.grad. The new W is already a different, non-leaf tensor, and its .grad is None.
• TypeError: unsupported operand type(s) for *: 'float' and 'NoneType' — you never called loss.backward(). W.grad is still None.
• three FAILs at once: loss is nan at step 61, accuracy after 100 steps is 9.4%, loss 2.3207 -> nan — you forgot zero_(). The loop looks like it works — and falls over only after a few dozen steps. That is what step 6 is about.