pytorch 计算模型的参数量

参考于：https://blog.csdn.net/jdzwanghao/article/details/84196239

def model_structure(model):
    blank = ' '
    print('-'*90)
    print('|'+' '*11+'weight name'+' '*10+'|' \
            +' '*15+'weight shape'+' '*15+'|' \
            +' '*3+'number'+' '*3+'|')
    print('-'*90)
    num_para = 0
    type_size = 1  ##如果是浮点数就是4
    
    for index, (key, w_variable) in enumerate(model.named_parameters()):
        if len(key) <= 30: 
            key = key + (30-len(key)) * blank
        shape = str(w_variable.shape)
        if len(shape) <= 40:
            shape = shape + (40-len(shape)) * blank
        each_para = 1
        for k in w_variable.shape:
            each_para *= k
        num_para += each_para
        str_num = str(each_para)
        if len(str_num) <= 10:
            str_num = str_num + (10-len(str_num)) * blank
    
        print('| {} | {} | {} |'.format(key, shape, str_num))
    print('-'*90)
    print('The total number of parameters: ' + str(num_para))
    print('The parameters of Model {}: {:4f}M'.format(model._get_name(), num_para * type_size / 1000 / 1000))
    print('-'*90)

结果：

------------------------------------------------------------------------------------------
|           weight name          |               weight shape               |   number   |
------------------------------------------------------------------------------------------
| embed_in.0.weight              | torch.Size([28, 1, 1, 5, 5])             | 700        |
| embed_in.0.bias                | torch.Size([28])                         | 28         |
| downC.0.block1.0.weight        | torch.Size([28, 28, 1, 3, 3])            | 7056       |
| downC.0.block1.1.weight        | torch.Size([28])                         | 28         |
| downC.0.block1.1.bias          | torch.Size([28])                         | 28         |
| downC.0.block2.0.weight        | torch.Size([28, 28, 3, 3, 3])            | 21168      |
| downC.0.block2.1.weight        | torch.Size([28])                         | 28         |
| downC.0.block2.1.bias          | torch.Size([28])                         | 28         |
| downC.0.block2.3.weight        | torch.Size([28, 28, 3, 3, 3])            | 21168      |
| downC.0.block3.weight          | torch.Size([28])                         | 28         |
| downC.0.block3.bias            | torch.Size([28])                         | 28         |
| downC.1.block1.0.weight        | torch.Size([36, 28, 1, 3, 3])            | 9072       |
| downC.1.block1.1.weight        | torch.Size([36])                         | 36         |
| downC.1.block1.1.bias          | torch.Size([36])                         | 36         |
| downC.1.block2.0.weight        | torch.Size([36, 36, 3, 3, 3])            | 34992      |
| downC.1.block2.1.weight        | torch.Size([36])                         | 36         |
| downC.1.block2.1.bias          | torch.Size([36])                         | 36         |
| downC.1.block2.3.weight        | torch.Size([36, 36, 3, 3, 3])            | 34992      |
| downC.1.block3.weight          | torch.Size([36])                         | 36         |
| downC.1.block3.bias            | torch.Size([36])                         | 36         |
| downC.2.block1.0.weight        | torch.Size([48, 36, 1, 3, 3])            | 15552      |
| downC.2.block1.1.weight        | torch.Size([48])                         | 48         |
| downC.2.block1.1.bias          | torch.Size([48])                         | 48         |
| downC.2.block2.0.weight        | torch.Size([48, 48, 3, 3, 3])            | 62208      |
| downC.2.block2.1.weight        | torch.Size([48])                         | 48         |
| downC.2.block2.1.bias          | torch.Size([48])                         | 48         |
| downC.2.block2.3.weight        | torch.Size([48, 48, 3, 3, 3])            | 62208      |
| downC.2.block3.weight          | torch.Size([48])                         | 48         |
| downC.2.block3.bias            | torch.Size([48])                         | 48         |
| center.block1.0.weight         | torch.Size([64, 48, 1, 3, 3])            | 27648      |
| center.block1.1.weight         | torch.Size([64])                         | 64         |
| center.block1.1.bias           | torch.Size([64])                         | 64         |
| center.block2.0.weight         | torch.Size([64, 64, 3, 3, 3])            | 110592     |
| center.block2.1.weight         | torch.Size([64])                         | 64         |
| center.block2.1.bias           | torch.Size([64])                         | 64         |
| center.block2.3.weight         | torch.Size([64, 64, 3, 3, 3])            | 110592     |
| center.block3.weight           | torch.Size([64])                         | 64         |
| center.block3.bias             | torch.Size([64])                         | 64         |
| upS.0.1.weight                 | torch.Size([48, 64, 1, 1, 1])            | 3072       |
| upS.0.1.bias                   | torch.Size([48])                         | 48         |
| upS.1.1.weight                 | torch.Size([36, 48, 1, 1, 1])            | 1728       |
| upS.1.1.bias                   | torch.Size([36])                         | 36         |
| upS.2.1.weight                 | torch.Size([28, 36, 1, 1, 1])            | 1008       |
| upS.2.1.bias                   | torch.Size([28])                         | 28         |
| upC.0.block1.0.weight          | torch.Size([48, 48, 1, 3, 3])            | 20736      |
| upC.0.block1.1.weight          | torch.Size([48])                         | 48         |
| upC.0.block1.1.bias            | torch.Size([48])                         | 48         |
| upC.0.block2.0.weight          | torch.Size([48, 48, 3, 3, 3])            | 62208      |
| upC.0.block2.1.weight          | torch.Size([48])                         | 48         |
| upC.0.block2.1.bias            | torch.Size([48])                         | 48         |
| upC.0.block2.3.weight          | torch.Size([48, 48, 3, 3, 3])            | 62208      |
| upC.0.block3.weight            | torch.Size([48])                         | 48         |
| upC.0.block3.bias              | torch.Size([48])                         | 48         |
| upC.1.block1.0.weight          | torch.Size([36, 36, 1, 3, 3])            | 11664      |
| upC.1.block1.1.weight          | torch.Size([36])                         | 36         |
| upC.1.block1.1.bias            | torch.Size([36])                         | 36         |
| upC.1.block2.0.weight          | torch.Size([36, 36, 3, 3, 3])            | 34992      |
| upC.1.block2.1.weight          | torch.Size([36])                         | 36         |
| upC.1.block2.1.bias            | torch.Size([36])                         | 36         |
| upC.1.block2.3.weight          | torch.Size([36, 36, 3, 3, 3])            | 34992      |
| upC.1.block3.weight            | torch.Size([36])                         | 36         |
| upC.1.block3.bias              | torch.Size([36])                         | 36         |
| upC.2.block1.0.weight          | torch.Size([28, 28, 1, 3, 3])            | 7056       |
| upC.2.block1.1.weight          | torch.Size([28])                         | 28         |
| upC.2.block1.1.bias            | torch.Size([28])                         | 28         |
| upC.2.block2.0.weight          | torch.Size([28, 28, 3, 3, 3])            | 21168      |
| upC.2.block2.1.weight          | torch.Size([28])                         | 28         |
| upC.2.block2.1.bias            | torch.Size([28])                         | 28         |
| upC.2.block2.3.weight          | torch.Size([28, 28, 3, 3, 3])            | 21168      |
| upC.2.block3.weight            | torch.Size([28])                         | 28         |
| upC.2.block3.bias              | torch.Size([28])                         | 28         |
| embed_out.0.weight             | torch.Size([28, 28, 1, 5, 5])            | 19600      |
| embed_out.0.bias               | torch.Size([28])                         | 28         |
| out_affs_2.0.weight            | torch.Size([3, 28, 1, 1, 1])             | 84         |
| out_affs_2.0.bias              | torch.Size([3])                          | 3          |
------------------------------------------------------------------------------------------
The total number of parameters: 821531
The parameters of Model UNet_PNI: 0.821531M
------------------------------------------------------------------------------------------