跳转到主要内容

随着 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 的兼容性。Profiler 使用一个新的 GPU 分析引擎,该引擎使用 Nvidia CUPTI API 构建,能够高保真地捕获 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上提交问题。