检查PyTorch模型中的参数总数

Question

如何计算PyTorch模型中的参数总数？类似于model.count_params()Keras的东西.

Answer 1

PyTorch没有像Keras那样计算参数总数的函数,但是可以对每个参数组的元素数量求和:

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

如果您只想计算可训练的参数:

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

答受此启发答案上PyTorch论坛.

注意:我正在回答我自己的问题.如果有人有更好的解决方案,请与我们分享.

Answer 2

为了像 Keras 一样获取每一层的参数计数，PyTorch 有model.named_pa​​ramters()，它返回参数名称和参数本身的迭代器。

下面是一个例子：

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)

输出看起来像这样：

+-------------------+------------+
|      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

Answer 3

如果你想在不实例化模型的情况下计算每层的权重和偏差的数量，你可以简单地加载原始文件并迭代结果，collections.OrderedDict如下所示：

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)))

你会得到类似的东西

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

Answer 4

如果要避免重复计算共享参数，可以使用torch.Tensor.data_ptr. 例如：

sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())

这是一个更详细的实现，其中包括一个过滤不可训练参数的选项：

def numel(m: torch.nn.Module, only_trainable: bool = False):
    """
    returns the total number of parameters used by `m` (only counting
    shared parameters once); if `only_trainable` is True, then only
    includes parameters with `requires_grad = True`
    """
    parameters = list(m.parameters())
    if only_trainable:
        parameters = [p for p in parameters if p.requires_grad]
    unique = {p.data_ptr(): p for p in parameters}.values()
    return sum(p.numel() for p in unique)

Answer 5

你可以torchsummary用来做同样的事情。这只是两行代码。

from torchsummary import summary

print(summary(model, (input_shape)))