PyTorch Profiler输出解读及torchprof等效功能咨询

一、Name列含义解析 #

• 模型层名称：你看到的模型层名称对应PyTorch模型中定义的Module（比如自定义的卷积块、LSTM层等），和你代码里的模块结构一一对应。
• aten:xxx 前缀的含义：aten:开头的是PyTorch底层ATen库的核心运算操作。ATen是PyTorch的张量运算基础库，所有高层API最终都会调用这些底层算子：
  • aten:gru：对应GRU层的底层张量运算实现
  • aten:uniform：对应生成均匀分布张量的底层操作
    这些是模型层执行时实际跑的计算单元，是高层模块的具体执行逻辑拆解。

二、实现类似torchprof的模块耗时细分（替代prof.display(show_events=True)） #

要获得按模块层级拆分的耗时统计，可通过以下方式实现：

1. 初始化Profiler时开启必要参数：

   with torch.profiler.profile(
       activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
       record_shapes=True,
       with_stack=True,
       on_trace_ready=torch.profiler.tensorboard_trace_handler('./profiler_log')
   ) as prof:
       # 运行你的模型推理/训练代码
       model(input_tensor)
       prof.step()
   

2. 用TensorBoard查看层级耗时：
   启动TensorBoard后，进入Profile面板的Operator View，通过Group by选项选择Module，就能看到每个模块下的子模块及底层算子的耗时细分，逻辑和torchprof一致。
3. 终端输出层级耗时汇总：
   执行prof.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=100)，group_by_stack_n会按调用栈分组，近似实现模块层级的耗时汇总，不过可视化方式更直观。

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