一种简单的统计pytorch模型参数量的方法

Posted 2023-03-03 午夜零时

nelememt()函数

Tensor.nelement()->引自Tensor.numel()->引自torch.numel(input)
三者的作用是相同的

Returns the total number of elements in the input tensor.
返回当前tensor的元素数量

利用上面的函数刚好可以统计模型的参数数量

parameters()函数

Module. parameters (recurse=True)

Returns an iterator over module parameters.
recurse ( bool ) – if True , then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.
返回module的参数迭代器，这里的参数parameters其实就是tensor，在module中充当参数的作用
如果recurse为True：返回当前module和子module的参数
如果recurse为False：只返回当前module的成员参数，不返回子module的参数

统计module的参数量

from torchvision.models import resnet50
model = resnet50()
total = sum([param.nelement() for param in model.parameters()])
print("total = ", total)

输出：
total = 25557032