6.4 Profiler：补给线体检

6.4 Profiler：补给线体检

本节摘要：训练慢的时候，直觉常常指错方向。Profiler 给出一个 step 的完整耗时解剖：CPU 调度、GPU 计算、数据装载各占多少。本节做两次体检——一次找出算子热点，一次诊断数据瓶颈——建立"先测量、再动手"的性能优化流程。

时间的去向要靠测量

后勤第二站。训练比预期慢时，常见的第一反应是"换更大的显卡"，而体检往往会给出便宜得多的答案：瓶颈可能在 DataLoader 的单线程装车、在频繁的 CPU 与 GPU 之间搬运、甚至在你随手写的某个逐元素小算子上。不测量就优化，是在用预感分配预算。

PyTorch 自带的 Profiler 用法是"包住若干个 step，输出统计表"。它回答三个问题：时间都花在哪些算子上（热点排序）、GPU 真实在算还是闲着（占比）、CPU 与 GPU 之间的等待结构。前两者是本节重点，第三者进阶场景再查。

图 6-2：一个 step 的耗时解剖与对症下药

体检一：找出算子热点

用第 5 章的循环跑几步，看热点表：

import torch
import torch.nn as nn
from torch.profiler import profile, ProfilerActivity

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 10))
opt = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.CrossEntropyLoss()
x, y = torch.randn(64, 64), torch.randint(0, 10, (64,))

def one_step():
    opt.zero_grad()
    loss_fn(model(x), y).backward()
    opt.step()

one_step()                                   # 先空跑一步做预热
with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    for _ in range(10):
        one_step()
print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=5))
​

输出（版本不同列略有差异，结构一致）：

---------------------  ------------  ------------  ------------
                 Name     Self CPU     CPU total   Call Count
---------------------  ------------  ------------  ------------
          aten::linear       1.82ms      3.41ms           20
          aten::matmul       0.94ms      2.10ms           20
        aten::addmm        0.71ms      1.26ms           10
        aten::relu          0.42ms      0.88ms           10
        aten::backward      0.31ms      5.02ms           10
---------------------  ------------  ------------  ------------
​

读表的方法：Self CPU 是该算子自身开销（不含子调用），按它排序找热点；Call Count 暴露另一类问题——某个算子被调用几百次，多半是循环里有逐元素的小操作，批量化后往往一次就省下大半。CPU 侧的绝对时间在小模型上看着微小，放大到真实规模（层更多、batch 更大）就是等比放大的账。

体检二：诊断数据装载瓶颈

后勤瓶颈的经典再现：模型很轻、数据预处理很重，GPU 等 CPU 喂数据。Profiler 的 self_cpu_time_total 里数据相关算子占比高、或 GPU 利用率曲线呈锯齿状，都是症状。验证手段是开并行装车前后对比：

import time
import torch
from torch.utils.data import DataLoader, TensorDataset

# 模拟"预处理很重"的数据源：取样时做一段昂贵的计算
class HeavyDataset(TensorDataset):
    def __getitem__(self, i):
        img, label = super().__getitem__(i)
        for _ in range(3):                    # 假装这是一串预处理变换
            img = img * 0.999 + 0.001 * img.mean()
        return img, label

base = TensorDataset(torch.randn(512, 64), torch.zeros(512, dtype=torch.long))

for workers in [0, 4]:
    loader = DataLoader(HeavyDataset(base.tensors[0], base.tensors[1]),
                        batch_size=64, num_workers=workers)
    t0 = time.perf_counter()
    for _ in range(3):                        # 完整过三遍数据
        for xb, yb in loader:
            pass
    print(f"num_workers={workers}: {time.perf_counter()-t0:.2f}s")
​

输出（机器不同数值不同，关系稳定）：

num_workers=0: 1.94s
num_workers=4: 0.68s
​

结果：四条并行装车线把装载时间压到约三分之一。

解读：num_workers 的本质是"装车与计算流水线并行"——主进程训练时，子进程在旁边预先装下一批。配套参数 pin_memory=True（锁页内存）能让批次到 GPU 的搬运更快，两者几乎总是成对出现。并行数的甜蜜点与 CPU 核数相关，取"核数的一半到核数"之间试，超过核数反而互相抢资源。

完整案例：一次提速的完整决策链

背景：贯穿案例放大 32 倍（batch 2048、隐藏层 1024）后训练偏慢，预算有限只能做一项优化，用体检决定把钱花在哪。

操作：按顺序跑三组对照，每组只改一个变量：纯数据并行装载、混合精度、两者叠加。

# 实验记录（同一台机器、同一份数据、50 步计时）：
#   基线                       12.6s
#   + num_workers=4             9.8s   数据装载瓶颈被消除
#   + 混合精度（4.2 节）          7.1s   矩阵乘提速兑现
#   + 两者叠加                   5.4s   收益近似可加（瓶颈不同层）
print("决策：两项都做，先上并行装载（零成本），再上混合精度（改三行）")
​

输出：

决策：两项都做，先上并行装载（零成本），再上混合精度（改三行）
​

解读：这个假想但完全可复现的决策链示范了性能工作的正确姿势——每次只动一个变量，用同一把尺子（总时长）对比。如果反着来（一口气全开），你就永远不知道哪个改动值多少收益，下次换环境也无法决定优先级。体检的价值不在单次报告，而在它让优化从"试错"变成"决策"。

变式：在体检一里把 record_shapes=True 换成 with_stack=True 重跑，热点表会附带调用栈——当你发现某个想不到的算子在热点榜上时，调用栈直接告诉你是哪一行代码把它请来的。

本节要点回顾

• 先测量再优化：热点表排序看 Self 时间，调用次数暴露碎片化小操作；
• 三类瓶颈三副药：装载并行、热点合并、调度批量化；
• num_workers 与 pin_memory 成对，并行数不超 CPU 核数；
• 一次只动一个变量，每个改动都要有单独的对照组。

最后一站：模型练成了，怎么体面地交付。