pytorch 网络参数显存计算

一、使用torchstat

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

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

二、使用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：经过论坛咨询，报错的原因找到了，只需要把

```bash
pip install torchsummary

修改为

pip install torch-summary

三、计算参数量与可训练参数量

def get_parameter_number(model):
    total_num = sum(p.numel() for p in model.parameters())
    trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return {'Total': total_num, 'Trainable': trainable_num}

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