PyTorch Profiler 怎么导出可分享的 Chrome trace

想把 PyTorch 性能分析结果交给同事查看，核心不是截一张表，而是把 profiler 收集到的事件导出为 Chrome trace JSON。最小做法是用 torch.profiler.profile 包住要分析的迭代，在每轮调用 prof.step()，再通过 on_trace_ready 回调执行 prof.export_chrome_trace("trace.json")。

官方地址：https://pytorch.org/docs/stable/profiler.html

下面按一个可交付的小项目展开：输入是一个可重复调用的训练或推理函数，输出是一份可以复制给他人的 trace.json。示例只说明代码组织和导出位置，不虚构本机的耗时、算子排名或 GPU 结果。

项目目标：导出一份可复查的 trace 文件

先把目标收窄为三个结果：能明确哪些 activity 被采集，能控制采集发生在哪几轮，能把最后一个采集周期写成文件。这样排查“文件为空”“导出时机不对”时，有清晰的检查顺序。

准备最小 profiler 项目与采集范围

创建项目时把“业务工作负载”和“性能采集配置”分开。下面的 workload 只是占位函数；真实项目中可以换成一批训练、验证或推理代码。activities 只放当前环境支持且确实要观察的设备活动，CPU 是最小起点，CUDA 则应按运行环境决定是否加入。

from pathlib import Path

import torch


def workload(step: int) -> None:
    # 这里替换为一次真实的训练或推理迭代；不要把数据加载等无关阶段混进来
    x = torch.randn(256, 128)
    weight = torch.randn(128, 64)
    _ = x @ weight + step


trace_dir = Path("traces")
trace_dir.mkdir(parents=True, exist_ok=True)


def export_trace(prof: torch.profiler.profile) -> None:
    # 回调收到一个已完成采集周期的 profile；文件名要能区分周期或实验
    output = trace_dir / f"trace-step-{prof.step_num}.json"
    prof.export_chrome_trace(str(output))


activities = [torch.profiler.ProfilerActivity.CPU]
# 只有运行环境确实使用 CUDA 时，才将 CUDA activity 加入列表
if torch.cuda.is_available():
    activities.append(torch.profiler.ProfilerActivity.CUDA)

图中的 workload 是待分析对象，activities 声明采集范围，profile 保存事件，trace_dir 和 export_trace 负责形成文件。它们是职责边界，不表示某次运行必然包含某个设备或产生固定大小的文件。

用 schedule 和回调导出 Chrome trace

默认 profiler 可以连续记录整个上下文；当训练循环很长时，更适合用 schedule(wait=1, warmup=1, active=2, repeat=1) 把采集缩小为一个周期。等待阶段不记录，预热阶段让系统进入稳定状态，active 阶段才保存目标事件。每完成一次迭代都要调用 prof.step()，否则 schedule 不会按循环推进。

# schedule 的数字是示例参数，应该按任务长度和采集目的调整
schedule = torch.profiler.schedule(
    wait=1,    # 跳过首轮，避免把初始化阶段当成稳定样本
    warmup=1,  # 预热一轮，降低刚启动时的干扰
    active=2,  # 连续记录两轮，随后触发导出回调
    repeat=1,  # 只保存一个周期，便于先做小范围分享
)

with torch.profiler.profile(
    activities=activities,
    schedule=schedule,
    on_trace_ready=export_trace,
    record_shapes=True,  # 需要形状信息时开启；会增加采集开销
) as prof:
    for step in range(4):
        # 业务函数必须在 profiler 上下文中调用，才能进入采集范围
        workload(step)
        # 告诉 profiler 当前迭代结束，推进 wait/warmup/active 状态
        prof.step()

on_trace_ready 在采集周期达到保存条件时收到 profile 实例；回调里调用 export_chrome_trace 会写出 Chrome JSON 格式的 trace。官方文档也说明，启用形状或调用栈记录会增加额外开销，所以分享前应确认这些字段是否真的需要。

不要把 prof.step() 当成业务步骤编号：它只是向 profiler 发送“下一轮开始”的信号。回调拿到的是当前采集对象，导出方法会把采集到的 trace 写为 Chrome JSON；如果没有进入 active 周期，回调自然不会按预期产生文件。

打开、命名与分享导出文件

导出的文件名建议包含实验名、设备范围和周期编号，例如 bert-cpu-cuda-step-4.json。它是性能分析输入，不是模型 checkpoint，也不是 TensorBoard 日志目录；分享时应同时说明 PyTorch 版本、设备类型、输入规模和采集配置，避免别人把一份局部 trace 当成完整基准。

如果团队使用 Chrome trace 查看器，可把 JSON 文件交给该查看器分析事件轨道。本文不把查看器页面当成文章配图，也不为未执行的示例填写具体操作结果。你真正需要确认的是：导出文件能被接收方打开、文件名对应正确的实验、采集范围与问题一致。

常见坑与最终验收

最终验收可以按四个问题检查：是否明确了被分析函数；是否在上下文中推进了 prof.step()；是否配置了 on_trace_ready 和唯一文件名；是否记录了设备、版本和采集参数。满足这些条件后，导出文件才具备可复查的上下文。

为什么不能只调用一次 export_chrome_trace？因为导出方法只能写入已经采集到的事件。若没有正确配置上下文、周期和迭代信号，导出的内容可能为空或缺少目标阶段。

schedule 一定要用吗？不一定。短小任务可以让 profiler 连续记录整个上下文；长任务使用 schedule 更容易控制采集窗口，也更方便限制文件大小和额外开销。

Chrome trace 能直接证明哪一个算子最慢吗？它提供事件时间线和上下文，帮助你进一步分析；不能脱离设备、输入规模、采集配置和重复次数，单凭一份局部文件就下结论。

把 profiler 配置当成一个小项目交付：范围要小，周期要明，文件要可区分，说明要完整。这样生成的 Chrome trace 才能在团队之间复查，而不是只留下一份难以解释的 JSON。