使用torch.profiler时如何定位算子对应的Python调用代码行

问题场景 #

torch.profiler输出结果截图
已通过record_function为自定义逻辑标注名称，但无法定位aten::ne、SelectBackward这类底层算子对应的Python调用行；已知aten::ne对应torch.not_equals，但代码中无该算子的显式调用，需要精准定位每个profiler记录操作的实际调用位置。

具体实现方法 #

• 开启Profiler内置调用栈采集
  初始化profiler时传入with_stack=True是最直接的方案，开启后所有被采集的算子（含ATen底层算子、自动微分反向算子）都会绑定完整的Python调用栈信息，配置示例如下：

  import torch
  
  with torch.profiler.profile(
      # 按需选择采集CPU/CUDA活动
      activities=[
          torch.profiler.ProfilerActivity.CPU,
          torch.profiler.ProfilerActivity.CUDA
      ],
      with_stack=True, # 核心配置：开启调用栈追踪
      record_shapes=True, # 可选：同时记录张量输入形状，辅助定位
      profile_memory=True # 可选：按需开启显存/内存追踪
  ) as prof:
      # 放入需要分析的代码逻辑，比如训练步
      train_one_step(model, data)
  

  采集完成后有两种查看方式：

  1. 控制台打印：调用print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=-1))，输出的表格中每个算子条目都会附带对应的调用栈路径和行号
  2. 可视化查看：调用prof.export_chrome_trace("trace.json")导出Chrome trace格式，或导出到TensorBoard，点击对应算子事件即可展开完整调用链，直接跳转到对应Python代码行。

• 反向算子定位说明
  类似SelectBackward这类反向算子，是前向算子对应的自动求导节点，不会直接对应反向传播阶段你写的代码，开启栈追踪后，这类算子的调用栈会直接指向触发该反向节点的前向代码位置，顺着栈帧就能找到对应逻辑。

• 隐式调用算子定位说明
  很多ATen算子不会被显式调用，会被上层API甚至Python语法触发：比如aten::ne就会被张量的!=比较运算、布尔mask生成逻辑、部分框架内置分支判断（比如混合精度判断、DDP同步判断、loss计算内部的掩码逻辑）隐式触发，不需要手动在代码里搜torch.ne或torch.not_equals的显式调用，开启栈追踪后会直接显示触发该算子的最上层代码行。

注意：with_stack=True会带来5%~15%左右的性能开销，不需要全量训练周期开启，只在需要定位问题的迭代段开启即可，也可以搭配torch.profiler.schedule做分阶段采集，过滤掉数据加载、模型初始化等无关阶段的栈信息干扰。

内容的提问来源于stack exchange，提问作者HoliInn