5分钟快速上手 torchprof：pip 安装 + 完整示例，让 PyTorch 模型剖析变得简单-平芜编程栈

5分钟快速上手 torchprof：pip 安装 + 完整示例，让 PyTorch 模型剖析变得简单

【免费下载链接】torchprofPyTorch layer-by-layer model profiler项目地址: https://gitcode.com/gh_mirrors/to/torchprof

torchprof是一个专为 PyTorch 打造的逐层模型剖析工具（layer-by-layer model profiler）。只需一行pip install torchprof，你就能快速定位"模型里哪一层最耗时、哪一层最占显存"，是 PyTorch 模型性能分析与优化的轻量级好帮手。本文带你 5 分钟完成安装、运行完整示例、读懂剖析报告，并附常见问题解答。

🔍 为什么需要逐层剖析？

训练或推理变慢时，我们常会遇到这样的困惑：

模型整体耗时 2 秒，到底慢在卷积层还是全连接层？
显存不够用了，是哪一层在"偷吃"显存？

用time.time()手动计时只能拿到整体耗时；而 torchprof 会基于 PyTorch 内置的 autograd profiler，自动为每一层"挂钩"，按树状结构输出每个模块的耗时与内存指标，让瓶颈一目了然。

⚠️ 开始前的版本须知（30秒看完）

项目	说明
Python	3.5+
PyTorch	`>= 1.1.0, < 1.9`（见 setup.py 依赖声明）
许可证	MIT

💡重要提示：由于 PyTorch 1.9 对官方 profiler 做了大改动，torchprof 已标记为弃用（deprecated）。如果你使用的是 PyTorch 1.9 及以上版本，建议直接使用 PyTorch 官方内置的torch.profiler；但如果你维护的是 torch 1.1~1.8 的旧项目，torchprof 依然是最简单的逐层剖析方案。

📦 第一步：pip 安装（10秒）

pip install torchprof

依赖极轻——它只依赖 torch 本身，不引入任何额外重型库，装完即用。

🚀 第二步：完整示例（10行代码跑起来）

下面用 AlexNet 演示最标准的用法（CPU 环境即可运行）：

import torch import torchvision import torchprof model = torchvision.models.alexnet(pretrained=False) x = torch.rand([1, 3, 224, 224]) with torchprof.Profile(model) as prof: model(x) print(prof.display(show_events=False))

运行后你会得到一张模块树状报告（节选）：

Module | Self CPU total | CPU total | Number of Calls ---------------|----------------|-----------|---------------- AlexNet | | | ├── features | | | │├── 0 | 1.832ms | 7.264ms | 1 │├── 1 | 51.858us | 76.564us | 1 │├── 2 | 75.993us | 157.855us | 1 ... └── classifier | | | └── 6 | 62.258us | 77.356us | 1

📌有 GPU 怎么办？只需两处改动：模型和数据加.cuda()，剖析器加use_cuda=True：

model = torchvision.models.alexnet(pretrained=False).cuda() x = torch.rand([1, 3, 224, 224]).cuda() with torchprof.Profile(model, use_cuda=True, profile_memory=True) as prof: model(x)

此时报告会自动扩展出Self CUDA total、CUDA total以及 4 列显存指标，帮你区分"GPU 计算慢"还是"CPU↔GPU 传输慢"。

📊 第三步：读懂剖析报告的 5 个指标

列名	含义
Self CPU total	该层自身在 CPU 上花的时间（不含子模块）
CPU total	该层及其所有子层的 CPU 总耗时
Self CUDA / CUDA total	GPU 端对应指标（需`use_cuda=True`）
Self CPU Mem / CUDA Mem 等	该层占用/分配的内存（需`profile_memory=True`）
Number of Calls	该层前向被调用的次数

💡 概念小贴士：Self 时间= 这一层自己干的活；total 时间= 自己 + 所有子模块的总账。对比两者，就能找出"哪一层真正吃掉了时间"。

🔬 想看层内的底层算子？

传入show_events=True，即可下钻到每层内部的aten::conv2d、aten::relu_等底层操作：

print(prof.display(show_events=True))

这样就能回答"卷积慢在 cudnn 卷积核，还是慢在内存整理"这类更深的问题。

🎯 进阶：只剖析你关心的层

大模型逐层剖析开销较大？用paths参数精确指定要剖析的层（路径是元组形式的模块名路径，且不限于叶子层）：

paths = [("AlexNet", "features", "3"), ("AlexNet", "classifier")] with torchprof.Profile(model, paths=paths) as prof: model(x) print(prof)

其他层将被自动跳过，只输出你指定的层——调试时特别实用。

🧩 项目结构与源码导读

项目代码非常精简，核心就两个文件，适合通读学习：

文件	作用
`torchprof/profile.py`	核心`Profile`类：遍历模型子模块、为每层`forward`动态挂钩 profiler
`torchprof/display.py`	将剖析事件汇总成上面的树状表格报告
`torchprof/__init__.py`	对外导出`Profile`，当前版本 1.4.0
`tests/test_profile.py`	验证 AlexNet 各层的剖析事件结构
`tests/test_set_paths.py`	验证`paths`参数只剖析指定层
`setup.py`/`requirements.txt`	打包与依赖（torch 1.7.0 / torchvision 0.8.1 用于测试）

Profile的完整参数一览（定义于torchprof/profile.py）：

model：要剖析的模型
enabled=True：设为False可一键禁用剖析（方便调试开关）
use_cuda=False：是否统计 GPU 指标
profile_memory=False：是否统计内存（要求 torch ≥ 1.6）
paths=None：仅剖析指定路径的层

❓ 常见问题（FAQ）

Q1：报告里为什么只有 CPU 三列？A：默认只统计 CPU。GPU 列需use_cuda=True，内存列需profile_memory=True才会显示。

Q2：运行profile_memory=True时出现 RuntimeWarning？A：说明你的 PyTorch 版本低于 1.6，不支持内存剖析。torchprof 会自动降级为普通模式并继续运行，不影响结果。

Q3：剖析器可以反复使用吗？A：每个Profile实例在with块内不可重入（re-enter 会抛错），退出with后才能调用display()或raw()。不过你可以在同一个with块里多次调用model(x)，结果会自动聚合，Number of Calls会累加。

Q4：想要原始数据做二次分析？A：调用prof.raw()，会返回 PyTorch 原始的 EventList 数据，可自由导出、绘图或统计。