伴随 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上提交问题。