Check The Total Number Of Parameters In A PyTorch Model


Answer :

PyTorch doesn't have a function to calculate the total number of parameters as Keras does, but it's possible to sum the number of elements for every parameter group:

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

If you want to calculate only the trainable parameters:

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

Answer inspired by this answer on PyTorch Forums.

Note: I'm answering my own question. If anyone has a better solution, please share with us.


To get the parameter count of each layer like Keras, PyTorch has model.named_paramters() that returns an iterator of both the parameter name and the parameter itself.

Here is an example:

from prettytable import PrettyTable  def count_parameters(model):     table = PrettyTable(["Modules", "Parameters"])     total_params = 0     for name, parameter in model.named_parameters():         if not parameter.requires_grad: continue         param = parameter.numel()         table.add_row([name, param])         total_params+=param     print(table)     print(f"Total Trainable Params: {total_params}")     return total_params      count_parameters(net) 

The output would look something like this:

+-------------------+------------+ |      Modules      | Parameters | +-------------------+------------+ | embeddings.weight |   922866   | |    conv1.weight   |  1048576   | |     conv1.bias    |    1024    | |     bn1.weight    |    1024    | |      bn1.bias     |    1024    | |    conv2.weight   |  2097152   | |     conv2.bias    |    1024    | |     bn2.weight    |    1024    | |      bn2.bias     |    1024    | |    conv3.weight   |  2097152   | |     conv3.bias    |    1024    | |     bn3.weight    |    1024    | |      bn3.bias     |    1024    | |    lin1.weight    |  50331648  | |     lin1.bias     |    512     | |    lin2.weight    |   265728   | |     lin2.bias     |    519     | +-------------------+------------+ Total Trainable Params: 56773369 

If you want to calculate the number of weights and biases in each layer without instantiating the model, you can simply load the raw file and iterate over the resulting collections.OrderedDict like so:

import torch   tensor_dict = torch.load('model.dat', map_location='cpu') # OrderedDict tensor_list = list(tensor_dict.items()) for layer_tensor_name, tensor in tensor_list:     print('Layer {}: {} elements'.format(layer_tensor_name, torch.numel(tensor))) 

You'll get something like

conv1.weight: 312 conv1.bias: 26 batch_norm1.weight: 26 batch_norm1.bias: 26 batch_norm1.running_mean: 26 batch_norm1.running_var: 26 conv2.weight: 2340 conv2.bias: 10 batch_norm2.weight: 10 batch_norm2.bias: 10 batch_norm2.running_mean: 10 batch_norm2.running_var: 10 fcs.layers.0.weight: 135200 fcs.layers.0.bias: 260 fcs.layers.1.weight: 33800 fcs.layers.1.bias: 130 fcs.batch_norm_layers.0.weight: 260 fcs.batch_norm_layers.0.bias: 260 fcs.batch_norm_layers.0.running_mean: 260 fcs.batch_norm_layers.0.running_var: 260 

Comments

Popular posts from this blog

Are Regular VACUUM ANALYZE Still Recommended Under 9.1?

Can Feynman Diagrams Be Used To Represent Any Perturbation Theory?