随着 PyTorch 1.8.1的发布,一个全新改进的性能调试工具 PyTorch Profiler 来了。作为微软和 Facebook 合作的一部分,PyTorch Profiler 是一个开源工具,可以对大规模深度学习模型进行准确高效的性能分析和故障排除。
分析和改进大规模深度学习模型的性能是一个持续的挑战,随着模型规模的增加而变得越来越重要。长期以来,由于缺乏可用的工具,PyTorch 用户很难解决这个问题。市面上有标准的性能调试工具提供 GPU 硬件级别的信息,但是缺少了 PyTorch 特定操作的背景信息。为了恢复错过的信息,用户需要将多个工具组合在一起,或者手动添加最少的相关信息以使数据有意义。另外还有 autograd profiler(torch.autograd.profiler) ,它可以捕获 PyTorch 操作的信息,但不能得到详细的 GPU 硬件级信息,也不能提供可视化支持。
新的 PyTorch Profiler(torch.Profiler)是一个工具,它将这两种类型的信息结合在一起,然后构建实现这些信息全部潜力的经验。这个新的分析器收集 GPU 硬件和 PyTorch 相关信息,将它们关联起来,对模型中的瓶颈进行自动检测,并生成如何解决这些瓶颈的建议。来自 profiler 的所有信息都可以在 TensorBoard 中为用户可视化。新的 Profiler API 在 PyTorch 中得到了原生支持,并且提供了迄今为止最简单的体验,用户可以在不安装任何附加包的情况下分析他们的模型,并且可以通过新的 PyTorch Profiler 插件在 TensorBoard 中立即查看结果。下面是 PyTorch Profiler 的截图 —— 自动瓶颈检测。
PyTorch Profiler 是 PyTorch autograd profiler 的新一代版本。它有一个新的模块命名空间 torch.profiler,但保持了与 autograd profiler APIs 的兼容性。PyTorch Profiler 使用了一个新的 GPU 性能分析引擎,用 Nvidia CUPTI APIs 构建,能够高保真地捕获 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_trace=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 是 Python 开发人员和数据科学家中最流行的代码编辑器之一。最近,VS Code 的 Python 扩展将 TensorBoard 集成到代码编辑器中,包括对 PyTorch Profiler 的支持。一旦你安装了 VS Code 和 Python 扩展,就可以快速打开 TensorBoard Profiler 插件。使用快捷键 CTRL + SHIFT + P (Mac 上的 CMD + SHIFT + P) 启动命令面板,然后键入“Launch TensorBoard”命令即可。
这种集成带有内置的进程生命周期管理特性。VS Code 将自动安装 TensorBoard 包和 PyTorch Profiler 插件包(将于4月中旬推出) ,如果您的系统上没有这两个包的话。VS Code 还将为你启动 TensorBoard 进程,并自动在你的工作目录中查找任何 TensorBoard 日志文件。完成后,只需关闭选项卡,VS Code 就会自动关闭进程。再也不需要在系统中运行终端窗口来提供 TensorBoard UI 的后端!下面是运行在 TensorBoard 中的 PyTorch Profiler Trace View。
想了解更多 VS Code 中对 TensorBoard 的支持请查阅下面的博文:
https://devblogs.microsoft.com/python/python-in-visual-studio-code-february-2021-release/
参考: https://pytorch.org/blog/introducing-pytorch-profiler-the-new-and-improved-performance-tool/