pytorch获得模型的参数量和模型的大小

参考

buffer和parameter

在模型中，有buffer和parameter两种，其中parameter就是我们一般认为的模型的参数，它有梯度，可被训练更新。但是buffer没有梯度，不能被训练更新。
我们可以通过torch.nn.Module.buffers()和torch.nn.Module.named_buffers()返回模型中的buffer。第二个函数同时返回自己定义的名称和buffer。
同时，我们可以通过torch.nn.Module.parameters()和torch.nn.Module.named_parameters()返回模型中的parameter。第二个函数同时返回自己定义的名称和parameter。
两个函数都有一个bool型参数recurse，默认为true。如果为true，将递归的查找所有子层的参数。否则只查找第一层的子层。

torch.Tensor.nelement和torch.Tensor.element_size

我们得到的parameter和buffer都是Tensor类型的参数，而对于Tensor，第一个函数可以返回这个Tensor中的元素个数，比如矩阵中有多少数。第二个函数可以返回这个Tensor所对应的数据类型的字节大小。比如float32就是4字节。

获得模型的大小

def getModelSize(model):
    param_size = 0
    param_sum = 0
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()
        param_sum += param.nelement()
    buffer_size = 0
    buffer_sum = 0
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()
        buffer_sum += buffer.nelement()
    all_size = (param_size + buffer_size) / 1024 / 1024
    print('模型总大小为：{:.3f}MB'.format(all_size))
    return (param_size, param_sum, buffer_size, buffer_sum, all_size)

函数也很好理解，通过model.parameters()返回能迭代所有参数的迭代器，之后就能通过for循环得到所有的parameter。buffer也是类似。
返回的param_size是所有parameters的参数字节MB大小，buffer_size是所有buffer的参数字节MB大小，all_size就是模型的MB大小。

来源：eecspan