pytorch获得模型的参数信息，所占内存的大小

一 sum

一个模型所占的显存无非是这两种：

• 模型权重参数
• 模型所储存的中间变量

其实权重参数一般来说并不会占用很多的显存空间，主要占用显存空间的还是计算时产生的中间变量，当我们定义了一个model之后，我们可以通过以下代码简单计算出这个模型权重参数所占用的数据量：

import numpy as np

# model是我们在pytorch定义的神经网络层
# model.parameters()取出这个model所有的权重参数
para = sum([np.prod(list(p.size())) for p in model.parameters()])

# 下面的type_size是4，因为我们的参数是float32也就是4B，4个字节
 print('Model {} : params: {:4f}M'.format(model._get_name(), para * type_size / 1000 / 1000))

对上述含义的说明：https://oldpan.me/archives/how-to-use-memory-pytorch

二 torchsummary

1.pip install torchsummary安装

2.

import torch
from torchsummary import summary

# 需要使用device来指定网络在GPU还是CPU运行
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
netG_A2B = Generator(3, 3).to(device)   # 这里需要网络，Generator报错
summary(netG_A2B, input_size=(3, 256, 256))

Total params这三项比较好理解，因为有可能固定param。
input size也比较好理解：3*256*256/1024/1024*4=0.75(最后一个4表示存储是需要4字节，float32类型)
Params size也比较好计算：138,357,544/1024/1024*4=527.79
Forward/backward pass size (MB)的计算：(10*24*24+20*8*8+20*8*8+50+10)/1024/1024*4*2=0.064（注意最有还有个2）
https://blog.csdn.net/csdnxiekai/article/details/110517751

三  pytorch-model-summary

1.pip install pytorch-model-summary

2.

# show input shape
print(summary(Net(), torch.zeros((1, 1, 28, 28)), show_input=True))

# show output shape
print(summary(Net(), torch.zeros((1, 1, 28, 28)), show_input=False))

# show output shape and hierarchical view of net
print(summary(Net(), torch.zeros((1, 1, 28, 28)), show_input=False, show_hierarchical=True))

原文链接：https://blog.csdn.net/csdnxiekai/article/details/110517751