作者：Maxim Lukiyanov - 微软首席项目经理，Guoliang Hua - 微软首席工程经理，Geeta Chauhan - Facebook 合作伙伴工程主管，Gisle Dankel - Facebook 技术主管

伴随 PyTorch 1.8.1 版本的发布，我们很高兴宣布 PyTorch Profiler – 全新改进的 PyTorch 性能调试分析器。PyTorch Profiler 是微软和 Facebook 合作开发的开源工具，能够为大规模深度学习模型提供准确高效的性能分析和故障排除。

分析和改进大规模深度学习模型性能是一个持续存在的挑战，并且随着模型尺寸的增加，其重要性也日益增长。长期以来，由于缺乏可用工具，PyTorch 用户很难解决这一挑战。虽然有一些标准的性能调试工具可以提供 GPU 硬件级别的信息，但却遗漏了 PyTorch 特定的操作上下文。为了恢复遗漏的信息，用户需要将多个工具结合使用，或者手动添加最少的关联信息才能理解数据。此外，还有 autograd profiler (torch.autograd.profiler)，它可以捕获关于 PyTorch 操作的信息，但无法捕获详细的 GPU 硬件级别信息，也无法提供可视化支持。

全新的 PyTorch Profiler (torch.profiler) 是一款工具，它将这两种类型的信息结合在一起，并构建经验，充分发挥这些信息的潜力。这款新的分析器收集 GPU 硬件和 PyTorch 相关信息，将它们关联起来，自动检测模型中的瓶颈，并生成关于如何解决这些瓶颈的建议。来自分析器的所有这些信息都将在 TensorBoard 中可视化展示给用户。新的 Profiler API 在 PyTorch 中得到原生支持，并提供了迄今为止最简单的体验，用户无需安装任何额外的软件包即可分析他们的模型，并使用新的 PyTorch Profiler 插件在 TensorBoard 中立即查看结果。下图是 PyTorch Profiler 的屏幕截图 - 自动瓶颈检测。

开始使用

PyTorch Profiler 是 PyTorch autograd profiler 的下一代版本。它有一个新的模块命名空间 torch.profiler，但保持了与 autograd profiler API 的兼容性。Profiler 使用了一个新的 GPU 分析引擎，该引擎使用 Nvidia CUPTI API 构建，能够高保真地捕获 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) 启动命令面板，并输入 “Launch TensorBoard” 命令，快速打开 TensorBoard Profiler 插件。

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

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

反馈

请查看 PyTorch Profiler 文档，试用 Profiler 并告知我们您的体验。在 PyTorch 讨论论坛上提供您的反馈，或在 PyTorch GitHub 上提交问题。