如何使用一行代码获取PyTorch模型的参数大小

最近完善了一个Python库，做到了使用一行命令便可以获取到DNN网络的大小情况。现在进行发布，方便研究人员使用。

一、前言

为了获取到神经网络模型的feature map以及weight的大小，我尝试了很多方法，目前最方便的方法是用summary这个库（https://github.com/Pinging-ZJU/Pytorch-Memory-Utils），然而在我们的测试下发现这个库其实存在一些漏洞，使得对非线性网络的测量不准确（多测了很多层），于是我们对该库进行了修改，并添加了一些值方便用户阅读。

二、使用方法

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('\nEpoch: %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)
    ...
    ...

之后就可以跑代码了

三、结果

运行代码后得到如下打印：

image.png

即可以看到每一层的shape，大小（MB）以及weight的浮点数个数、大小（MB）信息。

GitHub库如下：

https://github.com/Pinging-ZJU/DNN-Printer

如果好用欢迎大家给个star！谢谢哦，这里也要感谢Peiyi的合作