pytorch

How to Print Number of Model Parameters in PyTorch

PyTorch is a popular framework that is used in deep learning. It offers multiple features for generating complex Neural Networks(NN). Users can perform model training operations with this framework. However, users need to be familiar with the number of parameters before training the model.

This blog will describe:

What are the Parameters in PyTorch?

In PyTorch, the “nn.Module” class is used for defining the models. It includes all the operations and layers that make up the model. Every layer contains a set of parameters. Parameters are basically updated while training to minimize the error between the model’s actual values and predictions.

Why do Users Need to Check the Parameters of the Model?

While training the model, users need to know about the number of parameters of their model because it takes a lot of memory and processing power. If they are familiar with the number of the model’s parameters, they can easily evaluate the amount of memory that will be required and how much time it will take to train which helps users optimize their training process as well as prevent the system from running out of space.

How to Display the Number of Model Parameters in PyTorch?

The “nn.Module” class has the “parameters()” method that is used to view the number of model parameters in the PyTorch model. To get all elements, the “num1()” method is used.

To understand the previously discussed concept, let’s have a look at the provided code:

import torch.nn as nn

class NNModel(nn.Module):
    def __init__(self):
        super(NNModel, self).__init__()
        self.fc1 = nn.Linear(10, 50)
        self.fc2 = nn.Linear(50, 1)

    def forward(self, i):
        i = self.fc1(i)
        i = self.fc2(i)
        return i

my_model = NNModel()
t_params = sum(p.numel() for p in my_model.parameters())
print(f"Total number of parameters: {t_params}")

In the above-stated code:

  • First, we define a model that has two linear layers.
  • Then, generate the model’s instance and utilize the “parameters()” method to retrieve all the parameters.
  • Next, we apply the generator expression to calculate all parameters by summing up each parameter’s number of elements.
  • Lastly, call the “print()” statement to display the resultant values on the screen:

In the above-described code, we have only displayed the total number of parameters, if you want to get the parameter’s name and size, the following lines of code can be used:

for name, param in my_model.state_dict().items():

print(name, param.size())

Here:

  • state_dict()” is the Python dictionary object that is utilized for storing and loading models from PyTorch.
  • item()” method is utilized to return the list with all dictionary keys along with values.
  • print()” statement is used to print the parameter’s name and size by passing the “size()” method and parameter:

That’s all! We have compiled the easiest way to print the number of model parameters in PyTorch.

Conclusion

In PyTorch, the “nn.Module” class is used for defining the models that include all the operations and layers that make up the model. The “nn.Module” class has the “parameters()” method that is used to view the number of model parameters in the PyTorch model. This write-up demonstrated the method for printing the number of model parameters in PyTorch.

About the author

Maria Naz

I hold a master's degree in computer science. I am passionate about my work, exploring new technologies, learning programming languages, and I love to share my knowledge with the world.