torch.utils.bottleneck

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 运算符就很有意义了。

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

警告

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

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