pytorch查看网络参数显存占用量等

转载自：https://blog.csdn.net/weixin_45292794/article/details/108227437，可能略有修改，本文之作个人纪录学习使用，版权归原作者所有。

1.使用torchstat

pip install torchstat


from torchstat import stat
import torchvision.models as models
model = models.resnet152()
stat(model, (3, 224, 224))

关于stat函数的参数，第一个应该是模型，第二个则是输入尺寸，3为通道数。我没有调研该函数的详细参数，也不知道为什么使用的时候并不提示相应的参数。

 

2.使用torchsummary

pip install torchsummary

from torchsummary import summary
summary(model.cuda(),input_size=(3,32,32),batch_size=-1)

使用该函数直接对参数进行提示，可以发现直接有显式输入batch_size的地方，我自己的感觉好像该函数更好一些。但是！！！不知道为什么，该函数在我的机器上一直报错！！！TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

 

Update：经过论坛咨询，报错的原因找到了，只需要把

pip install torchsummary

修改为

pip install torch-summary