torchstat

这是一个基于 PyTorch 的轻量级神经网络分析工具。 它旨在让网络的搭建快速简便，并具备调试能力。 注意：该仓库目前仍在开发中，因此部分 API 可能会发生变化。

此工具可以展示：

• 网络总参数量
• 理论浮点运算次数 (FLOPs)
• 理论乘加运算次数 (MAdd)
• 内存占用

安装

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

• 通过 pip 安装。

$ pip install torchstat

• 克隆本仓库后使用 setup.py 进行安装和更新。

$ python3 setup.py install

简单示例

如果你想尽快运行 torchstat，如果网络已存在于脚本中，可以直接将其作为命令行工具调用。否则，你需要将 torchstat 作为模块导入。

命令行工具

$ torchstat masato$ torchstat -f example.py -m Net
[MAdd]: Dropout2d 不支持！
[Flops]: Dropout2d 不支持！
[Memory]: Dropout2d 不支持！
      模块名称  输入形状 输出形状     参数量 内存(MB)           MAdd         Flops  MemRead(B)  MemWrite(B) duration[%]   MemR+W(B)
0           conv1    3 224 224   10 220 220      760.0       1.85   72,600,000.0  36,784,000.0    605152.0    1936000.0      57.49%   2541152.0
1           conv2   10 110 110   20 106 106     5020.0       0.86  112,360,000.0  56,404,720.0    504080.0     898880.0      26.62%   1402960.0
2      conv2_drop   20 106 106   20 106 106        0.0       0.86            0.0           0.0         0.0          0.0       4.09%         0.0
3             fc1        56180           50  2809050.0       0.00    5,617,950.0   2,809,000.0  11460920.0        200.0      11.58%  11461120.0
4             fc2           50           10      510.0       0.00          990.0         500.0      2240.0         40.0       0.22%      2280.0
total                                        2815340.0       3.56  190,578,940.0  95,998,220.0      2240.0         40.0     100.00%  15407512.0
===============================================================================================================================================
总参数量：2,815,340
-----------------------------------------------------------------------------------------------------------------------------------------------
总内存：3.56MB
总 MAdd：190.58MMAdd
总 Flops：96.0MFlops
总 MemR+W：14.69MB

如果你不确定如何使用某个特定命令，可以使用 -h 或 –help 选项来查看帮助信息，其中会列出该命令的用法及可用选项。

模块

from torchstat import stat
import torchvision.models as models

model = models.resnet18()
stat(model, (3, 224, 224))

功能与待办事项

注意：这些功能仅适用于 nn.Module。torch.nn.functional 中的模块目前尚不支持。

• FLOPs
• 参数数量
• 总内存
• Madd(FMA)
• MemRead
• MemWrite
• 模型摘要（详细信息，逐层）
• 导出评分表
• 任意输入形状

有关支持的层，请参阅 详细信息。

要求

• 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。

torchstat 快速上手指南

torchstat 是一个基于 PyTorch 的轻量级神经网络分析工具，旨在帮助开发者快速构建和调试网络模型。它可以统计模型的参数量、理论浮点运算量（FLOPs）、乘加运算量（MAdd）以及内存使用情况。

环境准备

在使用 torchstat 之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / macOS / Windows
• Python 版本：3.6 及以上
• 核心依赖：
  • PyTorch >= 0.4.0
  • Pandas >= 0.23.4
  • NumPy >= 1.14.3

提示：国内用户建议使用清华或阿里镜像源加速依赖安装。

安装步骤

您可以通过以下两种方式之一安装 torchstat：

方式一：通过 pip 安装（推荐）

pip install torchstat -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：源码安装

如果您需要最新功能或进行开发，可以克隆仓库后安装：

git clone https://github.com/Swall0w/torchstat.git
cd torchstat
python3 setup.py install

基本使用

torchstat 支持两种使用模式：作为命令行工具（CLI）直接运行脚本，或作为 Python 模块导入使用。

模式一：命令行工具 (CLI)

如果您的网络模型定义在一个独立的 Python 脚本中（例如 example.py），且模型类名为 Net，可以直接在终端运行：

torchstat -f example.py -m Net

运行后将输出详细的层级分析报告，包含输入/输出形状、参数量、内存占用、MAdd、Flops 等数据，并在底部汇总总计值。

• 查看帮助：如需了解具体命令选项，可运行 torchstat -h。

模式二：Python 模块调用

您可以在代码中直接导入 stat 函数来分析已实例化的模型。以下是一个分析 ResNet18 的简单示例：

from torchstat import stat
import torchvision.models as models

# 加载模型
model = models.resnet18()

# 执行分析，第二个参数为输入张量的形状 (C, H, W)
stat(model, (3, 224, 224))

注意事项：

• 目前主要支持 nn.Module 定义的层，torch.nn.functional 中的部分操作可能不被完全支持（如 Dropout2d 会显示不支持提示）。
• 输入形状需根据模型实际需求调整（例如灰度图通常为 (1, 224, 224)）。