QBoard » Artificial Intelligence & ML » AI and ML - PyTorch » Check the total number of parameters in a PyTorch model

Check the total number of parameters in a PyTorch model

  • How to count the total number of parameters in a PyTorch model? Something similar to model.count_params() in Keras.
      October 20, 2021 1:15 PM IST
    0
  • 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
      October 29, 2021 3:15 PM IST
    0
  • 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.
      October 20, 2021 3:35 PM IST
    0
  • Straight and simple

    print(sum(p.numel() for p in model.parameters()))
    
      January 14, 2022 2:08 PM IST
    0