Count the Total Number of Parameters in PyTorch Model

When we are using pytorch to build an ai model, we may want to know how many parameters in this model. In this tutorial, we will use an example to show you how to do.

Here is an example:

batch_size = 32
W = 100
C = 80
se = SEModule(C)
size = sum(param.numel() for param in se.parameters()) / 1024 / 1024
print("Model parameter number %.2fMB" % size)

Here SEModule model is proposed in the tutorial:

Implement Squeeze-and-Excitation (SE) Block for 1D Matrix in PyTorch – PyTorch Tutorial

Run this code, we may see:

Model parameter number 0.01MB

Moreover, if you plan to see all trainable parameters in pytorch model, you can view:

List All Trainable Variables in PyTorch – PyTorch Tutorial