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

参考

• Finding model size
• Pytorch模型中的parameter与buffer
• What pytorch means by buffers?
• Pytorch中Module,Parameter和Buffer的区别
• torch.Tensor.element_size
• torch.Tensor.nelement

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字节。

获得模型的大小

defgetModelSize(model):
    param_size =0
    param_sum =0for param in model.parameters():
        param_size += param.nelement()* param.element_size()
        param_sum += param.nelement()
    buffer_size =0
    buffer_sum =0forbufferin model.buffers():
        buffer_size +=buffer.nelement()*buffer.element_size()
        buffer_sum +=buffer.nelement()
    all_size =(param_size + buffer_size)/1024/1024print('模型总大小为：{:.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大小。