PyTorch中统计模块参数的学习笔记

Posted 2022-07-12 songyuc

关于模型的参数统计，请参考博文《PyTorch几种情况下的参数数量统计·爽朗》；
这里我们参考了上面的博文，对不同的统计方式进行了测试，其结果是torchinfo的结果是最方便的[test_torch_parameters.ipynb]；
对于统计某个模块的参数，可以使用下面的代码：

sum(p.numel() for p in module.parameters() if p.requires_grad)

更好的方式是使用torchinfo，示例代码如下：

from torchinfo import summary

model = ConvNet()
summary(model, input_size=(16, 1, 28, 28))