如何使用一行代码获取PyTorch模型的参数大小
一、前言
为了获取到神经网络模型的feature map以及weight的大小，我尝试了很多方法，目前最方便的方法是用summary这个库（），然而在我们的测试下发现这个库其实存在一些漏洞，使得对非线性网络的测量不准确（多测了很多层），于是我们对该库进行了修改，并添加了一些值方便用户阅读。
二、使用方法
1 直接在服务器上调用：pip3 install DNN_printer
2 在代码中添加from DNN_printer import DNN_printer以及DNN_printer(net, (3, 32, 32),batch_size) 其中，第一个参数为net的定义，第二个(3,32,32)为输出数据的维度，第三个数据为batch size的大小。
即：
from DNN_printer import DNN_printer
batch_size = 512
def train(epoch):
print(
Epoch: %d % epoch)
net.train()
train_loss = 0
correct = 0
total = 0
// put the code here and you can get the result
DNN_printer(net, (3, 32, 32),batch_size)
...
...
之后就可以跑代码了
三、结果
运行代码后得到如下打印：
即可以看到每一层的shape，大小（MB）以及weight的浮点数个数、大小（MB）信息。
GitHub库如下：
转载链接：https://www.jianshu.com/p/09b0b5cbb684
