Modern Frameworks

You have now written a forward pass, a loss, a gradient and a training loop by hand. A framework does exactly those four jobs for you, on a GPU, for networks with billions of numbers. After this lesson you can read a framework's model definition and know what it costs, and your program counts a network's parameters the way a framework reports them.

What a framework automates

The loop inside every framework step is the five-line shape you already know: clear the old gradients, run the forward pass, compute the loss, run backward, let the optimizer step. The first line is the one people forget. Gradients accumulate into the parameters on every backward call, which is useful for splitting a big batch across several passes and a bug everywhere else.

Dynamic versus static graphs

PyTorch builds the computation graph while your code runs (define-by-run), so loops, if statements and debuggers just work. Early TensorFlow described the graph first and ran it later, which allowed optimisation but made debugging painful; TensorFlow 2 and Keras default to eager execution too. JAX takes a functional route: you write pure functions and ask for their gradient, their compiled form or their batched form. Different styles, same chain rule.

Counting parameters

Every framework can print a model's parameter count, and it is the first number to check. For a dense layer from a inputs to b outputs:

parameters = a*b weights + b biases

Take a network with layer sizes 100, 20, 5. Layer one has 10020 + 20 = 2,020 parameters. Layer two has 205 + 5 = 105. The total is 2,125. Activations like ReLU have none.

The count decides memory. Weights in 32-bit floats take 4 bytes each. Training with Adam also keeps a gradient and two moment estimates per parameter, so about 16 bytes per parameter before counting activations. A 7-billion-parameter model therefore needs roughly 112 GB of state just to train, which is why large models are split across many GPUs. These totals overflow a 32-bit integer long before that: a single 50,000 by 50,000 layer already holds 2.5 billion weights.

Hardware and precision

GPUs and tensor cores multiply big matrices far faster than a CPU. Mixed precision runs the forward and backward passes in 16-bit floats and keeps a 32-bit master copy of the weights for the update, which roughly halves memory traffic. Distributed training copies the model to every GPU and splits the batch (data parallel), or splits the model itself when it no longer fits on one device.

Common mistakes

• Forgetting the biases, or counting them per input instead of per output.
• Counting the input layer as if it had weights. Sizes are widths, and parameters live between consecutive pairs.
• Using a 32-bit integer for the total on large layers.
• Treating the parameter count as a speed number. Two models with equal counts can differ hugely in compute.

Your exercise

Each input line is a comma-separated list of layer sizes: input width, hidden widths, output width. The starter parses it and prints your number. Implement count_params so it returns the total number of weights and biases as a 64-bit integer. A list with fewer than two sizes has no layers and no parameters.