How to Get the Number of Parameters in PyTorch

The PyTorch framework for developing Machine Learning models uses a set of parameters to assess the performance of neural networks. Each parameter represents a transformation applied on the input to produce the output. In PyTorch, the class “nn.Module” manages all the parameters of a particular model. The training process for models involves the optimization of established parameters.

In this article, we will show how to check the number of parameters within a machine-learning model using Pytorch.

What are Parameters and Why Should You Know Their Number?

Machine learning model parameters are by themselves a subclass of “Tensors”. Parameters are defined as the “weights” and “biases” of a model that overlook the transformation of the input into the output. Both of these parameters are learnable and are optimized during the training of your model. However, weights are the primary parameters that are used to denote the linear transformation of data. On the other hand, biases are the supplemental parameters that define secondary data transformations.

How to Check the Number of Parameters in PyTorch?

Follow the code and steps given below to know the exact number of parameters of a neural network with a couple of linear layers in PyTorch:

Step 01: Install and Import Required Libraries

Use the “pip” installation package to install the “torch” library and the import method to import libraries into the project in Google Colab:

!pip install torch

import torch

import torch.nn as nn

Step 02: Define the Neural Network

Next, define the neural network model using the “nn.Module” class along with the name of the model. Then, use the “__init__” method to define the input and output sizes of the linear layers of the model as shown in the below code. Then, use the “forward” method to regulate the forward pass:

class Sample(nn.Module):

def __init__(self):
super(Sample, self).__init__()
self.fc1 = nn.Linear(1000, 750)
self.fc2 = nn.Linear(50, 10)

def forward(self, A):
A = torch.relu(self.fc1(A))
A = self.fc2(A)
return A

The above code works as follows:

Define the sample model with “nn.Module” class.
Use the “__init__” method to initialize this class.
Then, use the “nn.Linear” class to assign values to the variables.
Next, define the “forward” pass of the neural network.
Lastly, use the “torch.relu()” function to activate the layers of the neural network.

The screenshot is shown below:

Step 03: Instantiate the Model

Then, “instantiate” the sample model by creating a virtual instance and incorporating it into the definition of the variable “model”:

model = Sample()

Step 04: Calculate and Print the Number of Parameters

Lastly, count the number of parameters in the model via the “sum()” method and use the “print()” method to show the output as shown below:

total_params = sum(p.numel() for p in model.parameters())

print("Number of Parameters in the Sample Model: {:,}".format(total_params))

The output is shown below:

Note: You can access our Google Colab Notebook on how to check the number of parameters in PyTorch at this link.

Pro-Tip

The number of parameters in a model in PyTorch depends on its own complexity. A model with many different types of interconnected will have a tremendously large number of parameters that are brought into use during its training. Moreover, a large number of parameters in a model will also use greater processing power to deliver results.

Success! We have just shown how to process and showcase the number of parameters in a model in PyTorch.

Conclusion

To get the number of parameters in Pytorch, install and import PyTorch into the IDE. Then, use the “nn.Module” class to define the neural network. After that, utilize the “sum” method to calculate how many parameters are contained within a specific machine learning model on PyTorch. This blog has displayed how to access the number of parameters of a sample model.