Profiler 性能剖析

给模型做体检

模型能跑通之后，下一个问题往往是：为什么这么慢？哪个算子拖了后腿？是计算慢还是数据搬运慢？靠感觉猜不靠谱，得用工具量。

PyTorch 自带性能分析器（Profiler），能记录每个算子（operator）的耗时和内存占用。它就像体检报告，哪个指标异常，一眼就能看出来。注意别用旧接口 torch.autograd.profiler，那是遗留版本，PyTorch 1.8 起推荐的新 API 在 torch.profiler 里，本教程都按新 API 讲。

快速上手：一张耗时表格

最基本的用法，把要分析的代码放进 profile 上下文管理器：

import torch
from torch.profiler import profile, record_function, ProfilerActivity

model = torch.nn.Sequential(
    torch.nn.Linear(512, 1024),
    torch.nn.ReLU(),
    torch.nn.Linear(1024, 512),
)
x = torch.randn(32, 512)

with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    with record_function("forward"):
        y = model(x)

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=8))

跑出来是一张表，每行一个算子，带耗时、调用次数和输入形状。排序之后，最慢的算子就在最上面。profile 包住的代码会照常执行，分析不影响原逻辑，放心用。

读懂这张表

表里有两个容易混淆的列：CPU total 和 Self CPU。

打个比方，CPU total 是「一个项目连外包一起的总耗时」，Self CPU 是「自己亲手干的时间」。比如 aten::linear 内部会调用矩阵乘法，它的 CPU total 包含子算子，Self CPU 则扣掉了这部分。

所以定位瓶颈的正确姿势是：先按 CPU total 排序找大头，再按 Self CPU 看是不是这个算子本身慢。调用次数（# of Calls）也很重要，单个很快但被调用几万次的小算子，一样能把训练拖慢。

表格里还有什么

key_averages().table() 还带了几列关键信息：

• Self CPU %：该算子自耗时占总耗时的百分比，一眼看出谁是主角。
• CPU time avg：平均每次调用耗时。单次调用如果很慢，说明这个算子本身重。
• 想按输入形状进一步拆开，加 group_by_input_shape=True（前提是之前开了 record_shapes=True）。同一个算子处理不同形状的数据，耗时会分开显示。

给代码段贴标签

算子的名字像 aten::addmm，看不出对应哪段代码。用 record_function 给代码段起个名字，表格里就会出现你的标签：

# 沿用上一段定义的 model 和 x
with profile(activities=[ProfilerActivity.CPU]) as prof:
    with record_function("数据加载"):
        data = torch.randn(128, 512)
    with record_function("模型前向"):
        y = model(data)

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=8))

这样输出表里就有「数据加载」「模型前向」两行，各段耗时一目了然。训练循环里给每个环节贴一圈标签，瓶颈在哪段直接现形。

看内存和 GPU

加上 profile_memory=True，表格里会多出内存列，能看出每个算子分配了多少内存，谁在疯狂吃显存：

with profile(
    activities=[ProfilerActivity.CPU],
    profile_memory=True,
    record_shapes=True,
) as prof:
    y = model(x)

print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=8))

有 GPU 的话，把 ProfilerActivity.CUDA 加进 activities，就能看到每个算子在显卡上的真实耗时，排序键换成 cuda_time_total 即可。

追踪文件：看时间线

表格是汇总视图。想看每个算子发生的先后顺序，可以导出 Chrome 追踪文件：

prof.export_chrome_trace("trace.json")

打开 Chrome 浏览器，地址栏输入 chrome://tracing，把 trace.json 拖进去，就能看到一条算子时间线。CPU 和 GPU 的活动分别显示，谁在等谁、哪里有空档，一目了然。

长任务：只分析一小段

训练一跑几小时，全程记录不现实，文件会大到打不开。用 schedule 只采样一小段：

from torch.profiler import schedule

def trace_handler(p):
    print(p.key_averages().table(sort_by="self_cpu_time_total", row_limit=5))
    p.export_chrome_trace(f"trace_{p.step_num}.json")

with profile(
    activities=[ProfilerActivity.CPU],
    schedule=schedule(wait=1, warmup=1, active=2),
    on_trace_ready=trace_handler,
) as p:
    for step in range(8):
        y = model(x)
        p.step()

schedule 的含义是：先等 1 步、热身 1 步（记录但丢弃，消除启动抖动）、认真记录 2 步，然后重复。每轮记录结束，trace_handler 被调用，打印表格并存文件。p.step() 是告诉 profiler「这一步结束了」。

剖析完整训练步

光测前向不够，训练循环里 backward 和优化器更新也耗时。把它们一起包进 profile：

# 沿用上面的 model 和 x，以及 trace_handler 函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

with profile(
    activities=[ProfilerActivity.CPU],
    schedule=schedule(wait=2, warmup=1, active=3),
    on_trace_ready=trace_handler,
) as p:
    for epoch in range(3):
        for batch in range(20):
            y_pred = model(x)
            loss = (y_pred - torch.randn(32, 512)).pow(2).mean()
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            p.step()

表格里会多出 backward 相关的算子，哪个环节拖后腿，数据说话，不靠猜。

Timer：快速测时小工具

只想测某个操作的耗时，用不上 Profiler 这么重。torch.utils.benchmark.Timer 是更轻的选择：

from torch.utils.benchmark import Timer

timer = Timer(
    stmt="x @ y",
    setup="x = torch.randn(128, 128); y = torch.randn(128, 128)",
)
print(timer.blocked_autorange(min_run_time=0.5))

它会自动决定跑多少轮，输出中位数和四分位距，比手动 time.time() 掐表可靠得多。四分位距（IQR）表示多次测量的波动范围，波动越大说明这个操作越不稳定，可能是缓存、线程调度等外部因素在捣乱。

工具怎么选

性能工具不止一个，选对才省事：

• 手动 time.time()：只想看整段代码总耗时，最简单，够用。
• Timer：精确测单个操作，自动跑多轮排除波动。
• Profiler：找瓶颈、看算子明细、看内存，是主力工具。
• 还想把剖析结果和 TensorBoard 放一起？on_trace_ready 换成 torch.profiler.tensorboard_trace_handler("./logs")，TensorBoard 的 PROFILER 页面就能看。

由粗到细，各管一段。新手从 Timer 和 Profiler 入手就够。想对比改动前后的差距，固定用同一个排序键和 row_limit，两张表才有可比性。

一个真实的优化案例

官方教程里有个例子很有代表性。作者分析一个模块，发现 87% 的时间耗在一个叫 MASK INDICES 的代码段里。拆开看，问题出在两处：

1. 把 GPU 上的掩码张量拷回 CPU，就为了用 NumPy 的 argwhere，来回拷贝开销巨大。
2. 掩码用的是 double 精度，数据量白白翻倍。

修复也直接：全程留在 GPU 上，用 torch.nonzero 替代 NumPy；double 改成 float。结果总耗时从 5.9 秒降到 0.23 秒，快了二十多倍。

注意作者还做了一步：改完再跑一遍 Profiler，确认瓶颈真的消失了。优化的循环就是这样——量、改、再量，直到表格里没有刺眼的大头。

这就是 Profiler 的价值：不是玄学调参，是拿数据说话。

小结

性能剖析的思路就三步：profile 包住代码，表格看耗时排序，对症下药。配合 TensorBoard 和各类可视化，训练过程对你就是透明的了。还有一点提醒：优化别过度，先确认真的慢、慢在哪，再动手——Profiler 的表格就是判断依据。下一章，我们来收拾训练中最烦人的部分——各种报错与调试技巧。