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_shapesTrue, profile_memoryTrue, ) as prof: train_one_epoch()先采样一小段训练不要一上来 profile 全量。三、看表格也看时间线print(prof.key_averages().table(sort_bycuda_time_total, row_limit10))表格能告诉你哪些算子耗时最多时间线能看到 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、拷贝、算子、同步和优化器开销。先找瓶颈再谈技巧。性能优化靠证据不靠经验手感。