PyTorch Profiler 实战：先找瓶颈，再谈优化技巧

PyTorch Profiler 实战：先找瓶颈，再谈优化技巧

一、性能优化不能靠猜

PyTorch 训练慢时，常见反应是调 batch size、换显卡、开混合精度、改 DataLoader。问题是，瓶颈可能在数据读取、CPU 预处理、GPU kernel、通信同步或日志保存。没定位就优化，很容易忙半天没有效果。

Profiler 的价值，是把时间花在哪里说清楚。

二、先拆训练循环

flowchart TD A[DataLoader] --> B[CPU 预处理] B --> C[H2D 拷贝] C --> D[Forward] D --> E[Backward] E --> F[Optimizer Step]

训练吞吐由整条链路决定。GPU 利用率低，不一定是模型小，也可能是数据管线供不上。

import torch.profiler with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], record_shapes=True, profile_memory=True, ) as prof: train_one_epoch()

先采样一小段训练，不要一上来 profile 全量。

三、看表格也看时间线

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

表格能告诉你哪些算子耗时最多，时间线能看到 CPU 和 GPU 是否有空洞。两者结合，才能判断是算子慢还是等待多。

如果 GPU 时间线有大量空白，通常要检查 DataLoader、数据拷贝、同步点。tensor.item()、频繁打印、同步日志都可能偷偷让 CPU 等 GPU。

四、优化要做对照实验

Profiler 发现瓶颈后，每次只改一个变量。比如先提高num_workers，再测试pin_memory，再尝试混合精度。多个改动一起上，很难判断哪个有效。

profile_experiment: baseline_steps_per_sec: 42 change: enable_pin_memory new_steps_per_sec: 49

还要记录硬件、驱动、PyTorch 版本、batch size、输入长度。性能数据没有环境信息，后续很难复现。

Profiler 本身也有开销。不要在正式训练全程开启详细 profile。可以定期采样，或者只在性能回归时打开。

最后，优化目标要明确。训练吞吐、显存占用、收敛速度、端到端训练时长不是同一个指标。提高 steps/s 如果影响收敛质量，也未必划算。

Profiler 结果还要和数据规模匹配。小 batch 下发现的瓶颈，在大 batch 或分布式训练里可能变化。建议至少在目标 batch size 和目标输入长度上采样，否则优化方向可能偏。

profiler_run: warmup_steps: 10 active_steps: 20 batch_size: target input_shape: target

还要把 profile 产物保存下来。Chrome trace、表格摘要、环境信息和代码 commit 应该一起存进实验记录。这样别人才能复查，而不是只看到一句“DataLoader 是瓶颈”。

最后，Profiler 不应只在训练慢时使用。模型结构或数据管线大改后跑一次轻量 profile，可以提前发现性能回退。

五、总结

PyTorch Profiler 应先定位训练链路瓶颈，再逐项验证 DataLoader、拷贝、算子、同步和优化器开销。

先找瓶颈，再谈技巧。性能优化靠证据，不靠经验手感。