请点击上方“AI公园”，关注公众号

---

本文选自github

作者：Swall0w

编译：ronghuaiyang

---

【导读】

pytorch的工具，告诉你模型的统计信息，用了多少参数，浮点计算量多少，用了多少内存，做模型优化是非常有用！

github地址：https://github.com/lzhbrian/image-to-image-papers

这是一个基于Pytorch的轻量级的神经网络分析工具。它可以帮助你轻松愉快的创建你的神经网络，还可以进行调试。注意：这个代码仓还在开发中。所以，有些API以后可能会变。

这个工具可以显示：

• 总的神经网络的参数数量

• 理论上的浮点计算量 (FLOPs)

• 理论上的乘加操作总量 (MAdd)

• 内存的使用量

安装

有两种方法可以将torchstat安装到你的环境中。

• 使用pip安装

• 克隆这个仓库，然后使用setup的方法进行安装升级

简单的例子

如果你想立刻开始运行torchstat，如果你的模型在脚本里已经存在了，你可以像一个命令行工具一样调用它。否则你需要将torchstate作为一个模块导入进来。

命令行工具

如果你不知道怎么使用某个命令，你可以使用-h或者-help，然后你可以看到每个选项的帮助信息。

模块

功能& 待做

注意：这些功能只对 nn.Module有用， torch.nn.functional 中的模块还不支持。

• 
  

  浮点计算量（ FLOPs）

• 
  

  参数数量

• 
  

  总的内存消耗

• 
  

  乘加运算的总量(FMA)

• 
  

  读内存

• 
  

  写内存

• 
  

  模型的概要信息(细节, 每一层的)

• 
  

  导出评分表

• 
  

  任意的输入尺寸

所有支持的层：

Layer	Flops	Madd	MemRead	MemWrite
Conv2d	ok	ok	ok	ok
ConvTranspose2d		ok
BatchNorm2d	ok	ok	ok	ok
Linear	ok	ok	ok	ok
UpSample	ok
AvgPool2d	ok	ok	ok	ok
MaxPool2d	ok	ok	ok	ok
ReLU	ok	ok	ok	ok
ReLU6	ok	ok	ok	ok
LeaklyReLU	ok		ok	ok
PReLU	ok	ok	ok	ok
ELU	ok	ok	ok	ok

依赖

• Python 3.6+

• Pytorch 0.4.0+

• Pandas 0.23.4+

• NumPy 1.14.3+

参考

感谢 @sovrasov 的第一个版本的flops的计算, @ceykmc 的脚本的主干.

• flops-counter.pytorch

• pytorch_model_summary

• chainer_computational_cost

• convnet-burden.

本文可以任意转载，转载时请注明作者及原文地址。

请长按或扫描二维码关注我们