公司动态

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

📅 2026/8/27 17:04:20
5分钟快速上手 torchprof:pip 安装 + 完整示例,让 PyTorch 模型剖析变得简单
5分钟快速上手 torchprofpip 安装 完整示例让 PyTorch 模型剖析变得简单【免费下载链接】torchprofPyTorch layer-by-layer model profiler项目地址: https://gitcode.com/gh_mirrors/to/torchproftorchprof是一个专为 PyTorch 打造的逐层模型剖析工具layer-by-layer model profiler。只需一行pip install torchprof你就能快速定位模型里哪一层最耗时、哪一层最占显存是 PyTorch 模型性能分析与优化的轻量级好帮手。本文带你 5 分钟完成安装、运行完整示例、读懂剖析报告并附常见问题解答。 为什么需要逐层剖析训练或推理变慢时我们常会遇到这样的困惑模型整体耗时 2 秒到底慢在卷积层还是全连接层显存不够用了是哪一层在偷吃显存用time.time()手动计时只能拿到整体耗时而 torchprof 会基于 PyTorch 内置的 autograd profiler自动为每一层挂钩按树状结构输出每个模块的耗时与内存指标让瓶颈一目了然。⚠️ 开始前的版本须知30秒看完项目说明Python3.5PyTorch 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(pretrainedFalse) x torch.rand([1, 3, 224, 224]) with torchprof.Profile(model) as prof: model(x) print(prof.display(show_eventsFalse))运行后你会得到一张模块树状报告节选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_cudaTruemodel torchvision.models.alexnet(pretrainedFalse).cuda() x torch.rand([1, 3, 224, 224]).cuda() with torchprof.Profile(model, use_cudaTrue, profile_memoryTrue) as prof: model(x)此时报告会自动扩展出Self CUDA total、CUDA total以及 4 列显存指标帮你区分GPU 计算慢还是CPU↔GPU 传输慢。 第三步读懂剖析报告的 5 个指标列名含义Self CPU total该层自身在 CPU 上花的时间不含子模块CPU total该层及其所有子层的 CPU 总耗时Self CUDA / CUDA totalGPU 端对应指标需use_cudaTrueSelf CPU Mem / CUDA Mem 等该层占用/分配的内存需profile_memoryTrueNumber of Calls该层前向被调用的次数 概念小贴士Self 时间 这一层自己干的活total 时间 自己 所有子模块的总账。对比两者就能找出哪一层真正吃掉了时间。 想看层内的底层算子传入show_eventsTrue即可下钻到每层内部的aten::conv2d、aten::relu_等底层操作print(prof.display(show_eventsTrue))这样就能回答卷积慢在 cudnn 卷积核还是慢在内存整理这类更深的问题。 进阶只剖析你关心的层大模型逐层剖析开销较大用paths参数精确指定要剖析的层路径是元组形式的模块名路径且不限于叶子层paths [(AlexNet, features, 3), (AlexNet, classifier)] with torchprof.Profile(model, pathspaths) as prof: model(x) print(prof)其他层将被自动跳过只输出你指定的层——调试时特别实用。 项目结构与源码导读项目代码非常精简核心就两个文件适合通读学习文件作用torchprof/profile.py核心Profile类遍历模型子模块、为每层forward动态挂钩 profilertorchprof/display.py将剖析事件汇总成上面的树状表格报告torchprof/__init__.py对外导出Profile当前版本 1.4.0tests/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.pymodel要剖析的模型enabledTrue设为False可一键禁用剖析方便调试开关use_cudaFalse是否统计 GPU 指标profile_memoryFalse是否统计内存要求 torch ≥ 1.6pathsNone仅剖析指定路径的层❓ 常见问题FAQQ1报告里为什么只有 CPU 三列A默认只统计 CPU。GPU 列需use_cudaTrue内存列需profile_memoryTrue才会显示。Q2运行profile_memoryTrue时出现 RuntimeWarningA说明你的 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 数据可自由导出、绘图或统计。 小结你想知道对应操作哪层最耗时print(prof)看树状报告哪层占显存profile_memoryTrue层内算子细节show_eventsTrue只看某几层paths[...]从pip install torchprof到拿到第一张逐层剖析报告全程不到 5 分钟。对于 torch 1.1~1.8 环境下的 PyTorch 模型性能分析这是一个零配置、极简依赖的实用选择。快去找出你模型里最慢的那一层吧【免费下载链接】torchprofPyTorch layer-by-layer model profiler项目地址: https://gitcode.com/gh_mirrors/to/torchprof创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考