Inspecting Model Parameters

Once a neural network has been built, a natural first question is how big it is. Not how many layers it has or how many neurons each layer contains, but how many individual numbers the training process is actually allowed to adjust. That count is the total number of trainable parameters. It tells you how much the model can change to fit data, and it is one of the simplest checks you can run on a model you have just constructed.

This article follows a short path. You will read a two-line piece of code that performs the count, predict what each part contributes, and then check that reading against an explanation of each piece. After that, you will see where the trainable numbers live, what happens when a layer is frozen, and how a trained model is saved to a file with a single torch.save call.

The count is produced by a single line that uses sum() with a generator expression. Read it first, then predict what each part contributes before you read the explanation that follows.

The expression inside sum() has three parts, and each does a separate job. The first part, model.parameters(), is a sequence of parameter tensors that walks through every learnable tensor the model holds. The second part, the condition if p.requires_grad, throws away any tensor whose requires_grad attribute is False. The third part, p.numel(), turns each surviving tensor into a number: the count of individual elements inside it.

The formula and the Python code express the same rule. The code is simply the formula written with a generator expression. Because the sum adds element counts rather than tensor counts, the total answers the question of how many numbers can change. Each tensor's contribution is its numel(), and tensors that are not trainable contribute nothing.

In the multilayer networks of this chapter, the trainable parameters sit inside torch.nn.Linear layers. A Linear layer takes its inputs, multiplies them by a weight matrix, and adds a bias vector. Written compactly, its output is $y = Wx + b$, where $W$ is the weight matrix, $x$ is the input, and $b$ is the bias vector. Both $W$ and $b$ are stored as parameters of the layer, which is why both are counted.

Linear layers are sometimes called feedforward or fully connected layers. These names describe the same layer, not a different kind of parameter. Whatever name you use, the weight matrix and bias vector inside it are the trainable parameters that the count picks up.

The requires_grad attribute is the switch. A parameter does not need any special type or placement to be trainable. What makes it trainable is that its requires_grad flag is set to True. When the flag is False, the same tensor is still part of the model, but it is left out of the count and is not updated during training.

After a model has been trained, its learned parameters are the result of all that work. If those parameters exist only in the running Python session, they are lost when the session ends. The state dictionary is the tool that captures them. Calling model.state_dict() returns a dictionary containing the state of the model, including its parameters. On its own, it does not write a file.

Each entry in the dictionary pairs a string key with a tensor. The key is the name of one parameter, and the value is the tensor of numbers that holds its current learned values. For the two-layer example above, the first layer's weight matrix would appear under one key and its bias vector under another. The diagram shows how a key maps to its tensor for a small generated layer with a 2 by 3 weight matrix and a 2-element bias.

The source for this topic is a single line of code, run as the last line of a training script where model is a trained PyTorch model:

The data moves in a fixed order. The model holds its parameters. The call model.state_dict() reads those parameters and returns them as a dictionary. That dictionary is passed as the first argument to torch.save(), which writes it to the file named by the second argument. The model itself never appears as the thing being written. Only the dictionary produced by state_dict() is written.

When a save does not produce the result you expect, locate the failure by checking the three roles in order. First check the data: is the first argument a state dictionary produced by model.state_dict()? Second check the function: is the call to torch.save()? Third check the destination: is the second argument the file name you intended? Most failures sit in one of these three places, and the first check is where the result most often diverges.