由 Maxim Lukiyanov（微软高级项目经理）、Guoliang Hua（微软高级工程经理）、Geeta Chauhan（Facebook 合作伙伴工程主管）、Gisle Dankel（Facebook 技术主管）共同撰写

随着 PyTorch 1.8.1 版本的发布，我们激动地宣布推出 PyTorch Profiler——PyTorch 的新版性能调试分析器。PyTorch Profiler 是微软和 Facebook 合作开发的开源工具，能够为大规模深度学习模型提供准确高效的性能分析和故障排除。

分析并提高大规模深度学习模型性能是一个持续性的挑战，随着模型规模的增加，其重要性也在不断增长。长期以来，PyTorch 用户由于缺乏可用工具，在解决这一挑战时遇到了困难。虽然存在提供 GPU 硬件级别信息的标准性能调试工具，但它们缺少 PyTorch 特定操作的上下文。为了恢复丢失的信息，用户需要组合多个工具或手动添加最小相关性信息来理解数据。此外，还有自动微分分析器（ torch.autograd.profiler ），它可以捕获 PyTorch 操作的信息，但不能捕获详细的 GPU 硬件级别信息，也不能提供可视化支持。

新的 PyTorch 分析器（ torch.profiler ）是一个将两种类型的信息合并在一起并构建实现该信息全部潜力的经验的工具。这个新的分析器收集 GPU 硬件和 PyTorch 相关联的信息，将它们关联起来，自动检测模型中的瓶颈，并生成如何解决这些瓶颈的建议。分析器中的所有这些信息都通过 TensorBoard 为用户可视化。新的分析器 API 在 PyTorch 中原生支持，提供了迄今为止最简单的体验，用户可以在不安装任何额外包的情况下分析他们的模型，并立即在 TensorBoard 中使用新的 PyTorch 分析器插件看到结果。下面是 PyTorch 分析器——自动瓶颈检测的截图。

开始使用

PyTorch Profiler 是 PyTorch autograd profiler 的下一个版本。它有一个新的模块命名空间 torch.profiler ，但与 autograd profiler API 保持兼容。Profiler 使用一个基于 Nvidia CUPTI API 构建的新 GPU 分析引擎，能够以高保真度捕获 GPU 内核事件。要分析您的模型训练循环，请将代码包裹在 Profiler 上下文管理器中，如下所示。

 with torch.profiler.profile(
    schedule=torch.profiler.schedule(
        wait=2,
        warmup=2,
        active=6,
        repeat=1),
    on_trace_ready=tensorboard_trace_handler,
    with_stack=True
) as profiler:
    for step, data in enumerate(trainloader, 0):
        print("step:{}".format(step))
        inputs, labels = data[0].to(device=device), data[1].to(device=device)

        outputs = model(inputs)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        profiler.step()

schedule 参数允许您限制包含在分析中的训练步骤数量，以减少收集的数据量并简化通过关注重点来简化视觉分析。 tensorboard_trace_handler 自动将分析结果保存到磁盘，以便在 TensorBoard 中进行分析。

要在 TensorBoard 中查看分析会话的结果，请安装 PyTorch Profiler TensorBoard 插件包。

pip install torch_tb_profiler

Visual Studio Code 集成

Microsoft Visual Studio Code 是 Python 开发者和数据科学家最受欢迎的代码编辑器之一。VS Code 的 Python 扩展最近增加了 TensorBoard 的集成，包括对 PyTorch Profiler 的支持。一旦安装了 VS Code 和 Python 扩展，您可以通过使用快捷键 CTRL + SHIFT + P（在 Mac 上为 CMD + SHIFT + P）启动命令面板，并输入“启动 TensorBoard”命令来快速打开 TensorBoard Profiler 插件。

此集成自带生命周期管理功能。如果您系统上没有 TensorBoard 包和 PyTorch Profiler 插件包（将于 4 月中旬推出），VS Code 将会自动安装。VS Code 还会为您启动 TensorBoard 进程，并自动在当前目录中查找任何 TensorBoard 日志文件。完成操作后，只需关闭标签页，VS Code 将自动关闭进程。不再需要在您的系统上运行终端窗口来为 TensorBoard UI 提供后端！以下是 TensorBoard 中运行的 PyTorch Profiler Trace View。

在此博客中了解更多关于 VS Code 中 TensorBoard 支持的信息。

反馈

查看 PyTorch 分析器文档，尝试使用 Profiler，并分享您的体验。在 PyTorch 讨论论坛提供反馈或在 PyTorch GitHub 上提交问题。