Pytorch 学习笔记(21) : PyTorch Profiler

资料https://docs.pytorch.org/docs/stable/profiler.htmlhttps://docs.pytorch.org/tutorials/recipes/recipes/profiler_recipe.htmlhttps://github.com/ZhiqianXia/perf-compiler-learning/tree/main/08-pytorch/4-performance/profiler_labs 如果感兴趣可以下载玩玩一、概述PyTorch Profiler 是 PyTorch 提供的性能分析工具用于收集训练和推理过程中的性能指标。通过上下文管理器 API开发者可以识别最耗时的模型算子查看算子的输入形状和堆栈跟踪研究设备内核活动可视化执行跟踪⚠️注意torch.autograd模块中的旧版 API 已被弃用建议使用torch.profiler。二、核心 API 详解2.1 profile 上下文管理器推荐用法withtorch.profiler.profile(activities[torch.profiler.ProfilerActivity.CPU,torch.profiler.ProfilerActivity.CUDA,])asp:code_to_profile()print(p.key_averages().table(sort_byself_cuda_time_total,row_limit-1))主要参数说明参数类型说明activitiesiterable分析活动类型CPU、CUDA、XPUscheduleCallable调度器控制分析启停时机on_trace_readyCallable跟踪就绪时的回调函数record_shapesbool记录算子输入形状profile_memorybool跟踪张量内存分配/释放with_stackbool记录源代码信息文件和行号with_flopsbool估算特定算子的 FLOPSwith_modulesbool记录模块层次结构仅 TorchScriptacc_eventsbool跨多个分析周期累积事件post_processing_timeout_sfloat后处理超时时间2.2 调度器Schedule使用适用于长时间训练任务可在不同迭代获取多个跟踪deftrace_handler(prof):print(prof.key_averages().table(sort_byself_cuda_time_total,row_limit-1))withtorch.profiler.profile(activities[...],scheduletorch.profiler.schedule(wait1,warmup1,active2,repeat1),on_trace_readytrace_handler,)asp:foriterinrange(N):code_iteration_to_profile(iter)p.step()# 通知分析器新迭代开始schedule 参数说明wait等待步数warmup预热步数active活跃记录步数repeat重复周期数0 表示持续到结束skip_first跳过前 N 步skip_first_wait是否跳过首次等待2.3 TensorBoard 集成withtorch.profiler.profile(on_trace_readytorch.profiler.tensorboard_trace_handler(./log),...)asp:# 训练代码# 启动 TensorBoard# tensorboard --logdir ./log三、高级功能3.1 动态切换分析活动withtorch.profiler.profile(...)asp:code_to_profile_0()# 关闭 CUDA 活动收集p.toggle_collection_dynamic(False,[torch.profiler.ProfilerActivity.CUDA])code_to_profile_1()# 重新开启 CUDA 活动收集p.toggle_collection_dynamic(True,[torch.profiler.ProfilerActivity.CUDA])code_to_profile_2()3.2 Execution Trace Observer用于获取 AI/ML 工作负载的图表示支持重放基准测试fromtorch.profiler.execution_traceimportExecutionTraceObserverwithtorch.profiler.profile(execution_trace_observerExecutionTraceObserver().register_callback(./execution_trace.json),)asp:# 训练代码3.3 导出功能方法功能export_chrome_trace(path)导出 Chrome JSON 格式跟踪export_stacks(path, metric)保存堆栈跟踪key_averages()按算子名称分组平均事件events()获取未聚合的分析事件列表四、ProfilerActivity 类型CPU- CPU 活动CUDA- NVIDIA GPU 活动XPU- Intel GPU 活动MTIA- Meta Training and Inference AcceleratorHPU- Habana Gaudi 设备PrivateUse1- 私有自定义设备五、性能注意事项启用 shape 和 stack 追踪会产生额外开销record_shapesTrue时分析器会临时持有张量引用可能阻止某些优化并引入额外拷贝建议在调试时使用完整功能生产环境根据需求选择六、Intel ITT API可选针对 Intel 平台的额外支持torch.profiler.itt.is_available()# 检查 ITT 是否可用torch.profiler.itt.mark(msg)# 标记瞬时事件torch.profiler.itt.range_push(msg)# 压入嵌套范围torch.profiler.itt.range_pop()# 弹出嵌套范围七、完整示例代码importtorchimporttorch.profiler# 定义模型modeltorch.nn.Linear(10,10).cuda()inputtorch.randn(100,10).cuda()# 性能分析withtorch.profiler.profile(activities[torch.profiler.ProfilerActivity.CPU,torch.profiler.ProfilerActivity.CUDA,],scheduletorch.profiler.schedule(wait1,warmup1,active3,repeat2),on_trace_readytorch.profiler.tensorboard_trace_handler(./log),record_shapesTrue,profile_memoryTrue,with_stackTrue,with_flopsTrue,)asprof:forstepinrange(10):withtorch.profiler.record_function(model_forward):outputmodel(input)lossoutput.sum()withtorch.profiler.record_function(model_backward):loss.backward()prof.step()# 打印统计结果print(prof.key_averages().table(sort_bycuda_time_total,row_limit10))八、总结PyTorch Profiler 提供了从基础到高级的全方位性能分析能力场景推荐用法快速调试基础profile上下文管理器长期训练scheduleon_trace_ready可视化分析TensorBoard 集成内存优化profile_memoryTrue分布式训练tensorboard_trace_handler指定worker_name