torchprof 源码解析：动态 Hook forward 的 Monkey-Patching 剖析机制是如何实现的-北京尧图网络

torchprof 源码解析动态 Hook forward 的 Monkey-Patching 剖析机制是如何实现的【免费下载链接】torchprofPyTorch layer-by-layer model profiler项目地址: https://gitcode.com/gh_mirrors/to/torchproftorchprof 是一款轻量的 PyTorch 模型逐层性能剖析工具layer-by-layer profiler。它的核心思路是在模型运行时用Monkey-Patching猴子补丁技术动态替换每一层的forward方法套上 PyTorch 官方 autograd profiler从而输出每一层的 CPU/CUDA 耗时与显存/内存占用报表。本文完整拆解torchprof/profile.py中的实现机制带你看懂动态 Hook forward的全过程。⚠️先说结论torchprof 目前已处于废弃deprecated状态——PyTorch 1.9 重构官方 profiler 后它不再适用新项目请直接用torch.profiler。但它的补丁式动态 Hook手法非常经典对理解 PyTorch 模块机制和性能剖析原理极具价值。快速上手3 步完成 PyTorch 逐层性能剖析无需修改模型代码with语句包住一次前向传播即可核心用法不到 10 行import torch, torchvision, torchprof model torchvision.models.alexnet(pretrainedFalse) x torch.rand([1, 3, 224, 224]) with torchprof.Profile(model, use_cudaTrue, profile_memoryTrue) as prof: model(x) print(prof)输出是一张树形表格每个叶子层一行的耗时与内存节选自 README.md 示例此处为纯文本示意Module | Self CPU total | CPU total | Number of Calls ---------------|----------------|-----------|---------------- AlexNet | | | ├── features | | | │├── 0 | 1.832ms | 7.264ms | 1 │├── 1 | 51.858us | 76.564us | 1 ...想读源码先拉取仓库git clone https://link.gitcode.com/i/11bd72c7dcae3b7c6980eb157dfe3b3c整个库只有两个核心文件剖析逻辑在torchprof/profile.py表格渲染在torchprof/display.py测试用例在tests/test_profile.py——体量小非常适合精读。剖析机制全景查找 → 替换 → 恢复三步走torchprof 把动态 Hook拆成了清晰的生命周期正好对应 Python 的上下文管理器协议with torchprof.Profile(model) as prof: # 进入① 遍历模块 ② 替换 forward model(x) # 模型运行每个叶子层都被夹在 profiler 中测量 # 退出③ 恢复原始 forward下面逐步拆解。第一步递归遍历模块树找出每一层walk_modulesPyTorch 模型本质是一棵nn.Module树。torchprof/profile.py中的walk_modules用生成器深度优先遍历它产出一条条Trace(path, leaf, module)三元组——模块路径如(AlexNet, features, 0)、是否叶子层、模块实例本身def walk_modules(module, name, path()): named_children list(module.named_children()) path path (name,) yield Trace(path, len(named_children) 0, module) # 无子模块叶子层 for name, child in named_children: yield from walk_modules(child, namename, pathpath)叶子层leaf是默认剖析粒度不指定paths时只有叶子层会被 Hook父容器如features、classifier只是树形结构占位。第二步Monkey-Patching 核心——一行module.forward wrap_forward这是全库最精华的一行。进入上下文时Profile.__enter__torchprof 对每个被选中的模块执行替换_forward module.forward # 先保存原始 forward闭包捕获 self._forwards[path] _forward functools.wraps(_forward) def wrap_forward(*args, **kwargs): ... module.forward wrap_forward # ← 猴子补丁生效的唯一一行这里用到了 Python 属性查找的经典规则forward本是类属性定义在nn.Module或其子类上module.forward wrap_forward是把包装函数写入该实例的__dict__即实例属性之后访问module(...)触发module.forward时实例属性优先命中原始方法被遮蔽。这就是 Monkey-Patching 最直接的形式——不继承、不改动类定义、不调用任何 PyTorch 扩展 API纯靠 Python 属性机制完成动态 Hook。为什么不用官方register_forward_hookHook API 是观察者模式回调只能拿到输入/输出无法决定 profiler 何时开闭也无法包住整段执行。直接替换 forward 才能把 profiler 上下文精确地嵌进每次调用内部。几个工程细节值得注意均在torchprof/profile.py的_hook_trace中id(module)去重用内存地址集合防止同一子模块被引用多处时重复包装避免耗时被双重计入functools.wraps保留原方法的__name__等元信息避免调试时面目全非paths参数允许指定非叶子层精确剖析例如paths[(AlexNet, classifier)]其余层直接跳过。第三步wrap_forward 内部——在剖析上下文中运行原始 forward被替换后的 forward 长这样简化版def wrap_forward(*args, **kwargs): with torch_profiler.profile(use_cudaself.use_cuda) as prof: res _forward(*args, **kwargs) # 原始 forward 在剖析器夹缝中执行 self.trace_profile_events[path].append(prof.function_events) return res要点每层独立启停一次 autograd profiler事件列表按path归入trace_profile_events模型被调用多次时事件追加到同一列表这就是表格中 Number of Calls 列的来源版本兼容处理很典型torch 1.6 不支持profile_memory靠捕获TypeError降级重试并发出警告torch 1.8 需手动调用event_list.populate_cpu_children()补建树结构。第四步with 退出时自动恢复模型无痕还原def __exit__(self, exc_type, exc_val, exc_tb): tuple(map(self._remove_hook_trace, self.traces)) del self._forwards_remove_hook_trace把每个模块的forward从self._forwards里取回原函数并重新赋值模型对象在剖析结束后与从未被 Hook 过完全一致。此外__enter__中有重入检查二次进入会抛出RuntimeError(torchprof profiler is not reentrant)避免状态混乱。输出怎么读Self CPU total 与 CPU total 的区别渲染逻辑在torchprof/display.py每个模块聚合为一个Measure九元组CPU/CUDA 的 self 与 total 耗时、四类内存、调用次数。新手最容易混淆两列列名含义Self CPU total该层自己消耗的 CPU 时间不含子层/子算子CPU total含下层调用在内的累计时间例如 AlexNet 第一层卷积Self CPU total 1.832ms ≈ CPU total 7.264ms的差值就是它内部 conv、contiguous、add 等子算子链路的开销。加show_eventsTrue可以把每个aten::算子也展开成行定位到具体算子级别。⚠️ 避坑指南使用 torchprof 前必看已废弃setup.py 将 torch 限制在1.1.0,1.9PyTorch ≥ 1.9 请迁移到官方torch.profiler基于 Kineto能力更强。不可重入同一个Profile实例不能嵌套两次with。剖析有开销每个叶子层独立启停 profiler层数多的模型测量时间会明显变长用paths只剖析关键层可大幅提速。事件非确定性profiler 返回的算子顺序不保证稳定tests/test_profile.py中用子集包含断言而非顺序断言来验证 AlexNet 各层应出现的aten::算子。✅ 总结从 torchprof 学到的 3 个技巧实例属性遮蔽类属性是 Python 里最轻量的 Monkey-Patching 手段配合闭包保存原函数即可随时还原上下文管理器__enter__/__exit__是 Hook/Unhook 的天然配对结构保证异常路径下也能恢复现场用hasattr 异常捕获做跨版本兼容是维护依赖底层 API 工具时的实用套路。读懂了这套约 200 行的剖析核心你就能看懂大多数运行时插桩工具profiler、tracer、mock 框架的底层原理了。【免费下载链接】torchprofPyTorch layer-by-layer model profiler项目地址: https://gitcode.com/gh_mirrors/to/torchprof创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

torchprof 源码解析：动态 Hook forward 的 Monkey-Patching 剖析机制是如何实现的

相关推荐

性能优化揭秘：KISS-Matcher如何用TBB并行KD-Tree与Robin Hash实现百万点云的快速匹配

技术写作知识体系的构建方法

离线翻译怎么做？Argos Translate 四条命令完成本地机器翻译部署

更多资讯

电工证谁发的最管用在哪考河南电工证过期怎么处理全攻略

电工技师考电工证郑州报考避坑指南：急需特种作业证才能上岗找工作

东莞长安电工证报考指南：学历不高也能考，考前押题帮你稳过

电工证用不用年审吗？全国通用证书怎么报考不被坑

异地网上办理电工证吗？良心建议帮你搞懂转证流程

日照低压电工证好考么千万别踩坑这些细节你必须知道

今日更新

特种工指挥考试题库怎么刷才不会白交培训费？官方报名入口在这！

电工证考试多少分合格在哪考别再交冤枉钱了

福清应急管理局电工证真实经历：工地太忙根本没时间复习考试怎么办

本周精选

安监局电工证分数查询怎么查才能知道含金量高不高

新疆电工作业干活累吗良心建议助你避开考不过的坑

中牟县电工证怎么报名外省证能转吗

本月热门

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

读完文章，想聊聊您的网站？