作者: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) 是一个将这两种信息结合起来的工具,并在此基础上构建了能够充分发挥这些信息潜力的体验。这个新的 profiler 收集 GPU 硬件和 PyTorch 相关信息,并将它们关联起来,执行模型的自动瓶颈检测,并生成有关如何解决这些瓶颈的建议。所有这些 profiler 信息都在 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 提供后端支持了!下图是 PyTorch Profiler 的 Trace View 在 TensorBoard 中运行的截图。

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

反馈

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