0


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大小。


本文转载自: https://blog.csdn.net/qq_43219379/article/details/124003959
版权归原作者 eecspan 所有, 如有侵权,请联系我们删除。

“pytorch获得模型的参数量和模型的大小”的评论:

还没有评论