Neural networks end to end, and the PyTorch that runs them

Picture a mixing desk. Every input arrives on its own channel, and every channel has a fader that decides how much of it reaches the output. That is a neuron. The faders are called weights, and learning is nothing but moving them. So write it down. Multiply each input by its weight. Add them all up. Then add one more number called the bias, which shifts the whole thing up or down. That single result is the pre activation, and it is almost always written z. Then z goes through an activation function, and this is not decoration. Stack two linear layers with nothing between them and you get one linear layer, always. The bend is what buys you depth. Relu is the default: negatives become zero, positives pass straight through. Now put four of those neurons side by side, all reading the same three inputs. Each neuron has its own row of weights. Stack the rows and you have a matrix. One matrix multiply computes every neuron in the layer at once, which is the whole reason this runs fast. Shapes are what actually bite. Take sixty four flattened images of seven hundred and eighty four numbers each, times a matrix seven hundred and eighty four by one hundred and twenty eight. The result is sixty four by one hundred and twenty eight. The inner dimensions must match. That is the only rule. Stack those layers and you have a network. An input layer, some hidden layers, an output layer. Width is how many neurons sit in a layer. Depth is how many layers there are. The whole thing is one function: numbers in, numbers out, weights in between. Now follow one example the whole way through, with numbers small enough to check by hand. Two features going in. A hidden layer of three neurons. Two outputs, because there are two classes. Our input is two point zero, and minus one point zero. Start with the first hidden neuron. Its two weights are zero point five, and minus zero point two. Two times zero point five is one. Minus one times minus zero point two is zero point two. Add its bias of zero point one. That neuron's pre activation is one point three. Do the same for the other two, and the layer gives one point three, minus two point seven, and zero point seven. Relu clamps the negative one to zero. The hidden vector becomes one point three, zero, zero point seven. That middle neuron contributes nothing to this example. Second layer now, three weights per output. Multiply, sum, add the bias, and out come two numbers: one point zero eight, and minus zero point one one. Those are called logits. That is the entire forward pass, done by hand, on paper. Nothing changes for a batch. Instead of one row of features, you pass thirty two rows, and the same matrix multiply handles all of them at once. The bias is added to every row by broadcasting. One example or a thousand, the code is identical. In PyTorch that whole thing is three lines. A linear layer, a relu, another linear layer. The layer object holds its own weights and bias, and calling it runs the multiply. Notice what comes out: logits again, not probabilities. Those two numbers are not an answer yet. A logit is any real number at all. It can be negative, it can be enormous, and a pair of them will not add up to anything. Picture three judges holding up scores on their own private scales. Softmax fixes that in two steps. Exponentiate every logit, which makes them all positive. Then divide each one by the total. One point zero eight and minus zero point one one become zero point seven seven and zero point two three. They sum to one, and the ranking never changes. Now score the guess. Cross entropy takes the probability you gave the correct class, and returns minus its logarithm. Confident and right costs almost nothing. Confident and wrong costs enormously. Here the answer was class zero, so the loss is zero point two six. In PyTorch, cross entropy loss expects logits. It runs the softmax internally, in a numerically stable way. Hand it softmax output instead and you have applied softmax twice, the gradients go soft, and nothing crashes. It simply trains badly. If the output is a quantity rather than a class, use mean squared error instead. Subtract, square, average. Squaring punishes a big miss far harder than a small one. Classes get cross entropy, quantities get mean squared error, and choosing wrong fails quietly. Whichever you pick, the loss has to collapse to a single number. One scalar for the whole batch, averaged across it. That matters, because you can only ask how a small nudge changes one number. A vector has no single downhill direction. Here is the only question training ever asks. If I nudge this one weight up a tiny amount, does the loss go up or down, and by how much? That number is the partial derivative of the loss with respect to that weight. Its gradient. To answer it, PyTorch quietly records what you did. Every operation on a tensor adds a node to a graph: what went in, what came out, and how to differentiate itself. You wrote a forward pass. You also built a graph, for free. The graph is then walked backwards, using the chain rule. Each node receives the gradient of the loss with respect to its own output, multiplies that by its own local derivative, and hands the result to its inputs. Multiply along the chain. That is all backpropagation is. So calling backward starts at the loss with a gradient of one, and walks the graph in reverse. Every leaf tensor that asked for gradients ends up with a dot grad attribute, exactly the same shape as itself. One pass. Every weight in the network, at once. Two details bite everybody. Gradients only flow into tensors created with requires grad set to true. And backward adds to dot grad rather than replacing it, so if you forget to clear them, step two trains on step one's gradients as well. Then the update, which is one line of arithmetic. Move every weight a small step in the direction that lowers the loss. New weight equals old weight, minus the learning rate, times its gradient. Repeat a few hundred thousand times. That is training. PyTorch is four objects and five lines. The first object is the tensor: an array of numbers with a shape, a data type, and a device. Nearly every bug you will hit in your first month is one of those three being wrong. The second object is n n dot Module. Assign layers in the constructor and it registers them, so parameters collects every weight in the network however deeply it nests. Define forward, and calling the model runs it. That is the entire contract. Here it is for real. A sequential stack: seven hundred and eighty four numbers in, one hundred and twenty eight hidden, ten classes out. Notice there is no softmax at the end, and no backward method anywhere. Both of those are on purpose. The third object is the optimiser. Hand it the parameters once, and it owns the learning rate and any state it keeps. Plain gradient descent just steps downhill. Adam keeps a running average per parameter and adapts each step size. Adam is the default for a reason. And the fourth is the loop, which is five lines and never changes. Clear the gradients. Run the forward pass. Compute the loss. Call backward. Step the optimiser. Every model you have ever seen trained is those five lines, in that order. One more thing, at evaluation time. Call model dot eval, because dropout and batch norm behave differently outside training. Wrap it in no grad, so the graph is never built and the memory is never spent. Then switch back with model dot train. Picture a loading dock. Nothing is slow about lifting one box; the waste is walking back for the next one. That is why we train in batches. Sixty four examples in one trip costs far less than sixty four separate trips, and the gradient is steadier too. Devices next. A tensor lives either on the C P U or on the accelerator, and an operation needs both of its inputs in the same place. Move the model, and move the data. Forget one and you get the single most common error message in PyTorch. Now the surprise. The matrix multiply is almost never what is slow. Multiplying two n by n matrices reads n squared numbers, and does n cubed multiply-adds. So the bigger it gets, the more arithmetic you do per number you fetched. That shape of work runs beautifully. The slow part is moving numbers. An elementwise operation, like adding a bias or applying relu, does one cheap sum per element, but still reads the whole tensor out of memory and writes it back. Chain four of those and you have paid for four round trips. And upstream of all of it sits the data loader. If one process is decoding pictures while the accelerator waits, your expensive hardware is idle. Use several worker processes, pin the memory so the copy can overlap with compute, and the queue stays full. So measure before you optimise. The profiler tells you where the time actually went, and the answer is usually a copy, a wait, or a tiny kernel launched ten thousand times. Mixed precision and a fused compile step help. Guessing does not. So: a neuron is a weighted sum, a bias, and a bend. A layer is a matrix multiply. Softmax turns scores into probabilities, and a loss turns those into one number. Backward walks the graph and fills in every gradient. The loop is five lines. And the time goes into moving data, not multiplying it.