torch.utils.bottleneck

创建于：2018年3月23日 | 最后更新于：2022年9月28日

torch.utils.bottleneck 是一个工具，可以作为调试程序瓶颈的初步步骤。它使用 Python 剖析器和 PyTorch 的 autograd 剖析器来汇总脚本的运行情况。

在命令行上运行它：

python -m torch.utils.bottleneck /path/to/source/script.py [args]

其中 [args] 是传递给 script.py 的任意数量的参数，或者运行 python -m torch.utils.bottleneck -h 以获取更多使用说明。

警告

由于您的脚本将被剖析，请确保它能在有限的时间内退出。

警告

由于 CUDA 核函数的异步性质，在针对 CUDA 代码运行时，cProfile 输出和 CPU 模式的 autograd 剖析器可能无法显示正确的计时：报告的 CPU 时间表示启动核函数所花费的时间，但不包括核函数在 GPU 上执行的时间，除非该操作进行了同步。执行同步的操作在常规 CPU 模式剖析器下会显得异常昂贵。在这些计时不正确的场合，CUDA 模式的 autograd 剖析器可能很有帮助。

注意

要决定查看哪种（仅 CPU 模式或 CUDA 模式）autograd 剖析器输出，您应该首先检查您的脚本是 CPU 密集型（“CPU 总时间远大于 CUDA 总时间”）。如果是 CPU 密集型，查看 CPU 模式 autograd 剖析器的结果将有所帮助。另一方面，如果您的脚本大部分时间都花在 GPU 执行上，那么在 CUDA 模式 autograd 剖析器的输出中寻找负责的 CUDA 算子是有意义的。

当然，现实情况要复杂得多，您的脚本可能不会处于这两种极端情况之一，这取决于您评估的模型部分的哪个部分。如果剖析器的输出没有帮助，您可以尝试使用 nvprof 查看 torch.autograd.profiler.emit_nvtx() 的结果。但是，请注意 NVTX 的开销非常高，并且经常会给出严重失真的时间线。类似地，Intel® VTune™ Profiler 借助 torch.autograd.profiler.emit_itt() 进一步分析 Intel 平台的性能。

警告

如果您正在剖析 CUDA 代码，bottleneck 运行的第一个剖析器（cProfile）将在其时间报告中包含 CUDA 启动时间（CUDA 缓冲区分配成本）。如果您的瓶颈导致代码比 CUDA 启动时间慢得多，这应该不会有影响。

有关剖析器的更复杂用法（例如在多 GPU 的情况下），请参阅 https://docs.pythonlang.cn/3/library/profile.html 或 torch.autograd.profiler.profile() 以获取更多信息。