博客

介绍 PyTorch Profiler – 新改进的性能工具

随着 PyTorch 1.8.1 版本的发布，我们很高兴地推出 PyTorch Profiler——这是一款针对 PyTorch 全新改进的性能调试分析器。PyTorch Profiler 是微软与 Facebook 合作开发的开源工具，它能够对大规模深度学习模型进行准确、高效的性能分析与故障排除。

分析和改进大规模深度学习模型的性能是一个持续存在的挑战，并且随着模型规模的增长，其重要性日益凸显。长期以来，由于缺乏有效的工具，PyTorch 用户在解决这一问题时困难重重。现有的标准性能调试工具虽然能提供 GPU 硬件层面的信息，却缺失了 PyTorch 算子级别的上下文信息。为了获取这些缺失的信息，用户往往需要结合多种工具，或者手动添加最基本的关联信息来解读数据。此外，虽然之前有 autograd 分析器（torch.autograd.profiler）可以捕获 PyTorch 算子的信息，但它无法捕获详细的 GPU 硬件级别信息，也无法提供可视化支持。

全新的 PyTorch Profiler (torch.profiler) 将上述两类信息整合在一起，并构建了能够充分发挥这些数据价值的分析体验。这款新分析器既能收集 GPU 硬件信息，也能收集 PyTorch 相关信息，并将二者进行关联，自动检测模型中的性能瓶颈，并生成解决建议。所有这些来自分析器的数据都会在 TensorBoard 中为用户进行可视化展示。新的 Profiler API 已在 PyTorch 中原生支持，为用户提供了迄今为止最简便的使用体验：无需安装任何额外软件包即可对模型进行分析，并通过全新的 PyTorch Profiler 插件直接在 TensorBoard 中查看结果。下图展示了 PyTorch Profiler 的自动瓶颈检测功能。

入门指南

PyTorch Profiler 是 PyTorch autograd 分析器的下一代版本。它采用了全新的模块命名空间 torch.profiler，但仍保持与 autograd 分析器 API 的兼容性。该分析器使用基于 Nvidia CUPTI API 构建的全新 GPU 分析引擎，能够高保真地捕获 GPU 内核事件。要分析您的模型训练循环，请按照下方所示，将代码封装在分析器上下文管理器中。

 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 插件包（4 月中旬发布），VS Code 将自动为您安装。VS Code 还将为您启动 TensorBoard 进程，并自动在当前目录下查找任何 TensorBoard 日志文件。当您完成工作后，只需关闭标签页，VS Code 就会自动关闭相应的进程。从此无需再在终端窗口中运行进程来为 TensorBoard UI 提供后台支持！下图展示了在 TensorBoard 中运行的 PyTorch Profiler 追踪视图。

如需了解有关 VS Code 中 TensorBoard 支持的更多信息，请参阅这篇博客。

反馈

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