Calculate Number of Parameters in PyTorch Model

When working with PyTorch, knowing the number of parameters in the model is important for various reasons, such as model optimization, memory management, and performance evaluation. This knowledge helps to optimize the model's architecture and select appropriate hardware resources for efficient execution. This tutorial demonstrates how to calculate the number of parameters in a PyTorch model.

Code

In the following code, we define a simple model which consists of a linear layer with an input size of 100 and output size of 200, followed by a ReLU activation function, and finally, another linear layer with an input size of 200 and output size of 10.

To calculate the total number of parameters, we iterate over all the parameters of the model using parameters method and sum up the number of elements in each parameter tensor using numel method.

We're also calculating the number of trainable parameters. It achieved by adding the condition which checks requires_grad property. If the tensor has requires_grad set to true, then it can be updated during training.

import torch.nn as nn

model = nn.Sequential(
    nn.Linear(100, 200),
    nn.ReLU(),
    nn.Linear(200, 10)
)

total_params = sum(param.numel() for param in model.parameters())
print(total_params)

trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(trainable_params)

Here is the output you might expect from the code snippet:

22210
22210