新开源的 TensorFlow 性能分析工具 — TensorFlow Profiler

tfwechat · June 7, 2020, 2:59am

文 / Technical Writer, Anirudh Sriram 与产品经理 Gal Oshri，2020/4/15

性能是机器学习 (ML) 研究和生产解决方案成功与否的关键因素。提升模型训练的速度可加快迭代并减少开销。有时候，这也是实现特定机器学习解决方案的一项基本要求。
但是，要明确需要优化哪些部分却并非易事。我们需要思考是需要优化特定算子 (op)，亦或是输入流水线？
为帮助解答这类问题，我们开源了一套用于 TensorFlow 的综合性能分析工具——TensorFlow Profiler。这套工具除了能捕捉多方面的细节信息辅助分析，还能提供解决性能瓶颈（例如，输入受限程序）的一些建议。

底层专家可利用这些工具改进 TensorFlow 基础架构，Google 最热门产品的工程师们也使用它们来进行模型性能优化。我们希望广大社区能够充分利用我们所使用的这些工具，顺利完成性能分析。因此，我们最近开放了全新 TensorFlow Profiler 的源代码。

什么是 TensorFlow Profiler？

TensorFlow Profiler（以下简称 Profiler）提供了一套用于评估 TensorFlow 模型训练性能和资源消耗的工具。新版 Profiler 基于 Trace Viewer 等现有功能而构建，现已集成到 TensorBoard 中。

以下是 Profiler 最新提供的分析工具：

概述页：提供模型性能的顶级视图以及性能优化建议
输入流水线分析器：分析模型的数据输入流水线是否存在瓶颈，并为提升性能提出改进建议
TensorFlow 统计数据：显示分析会话期间执行的每个 TensorFlow 算子的性能统计数据
GPU 内核统计数据：显示每个 GPU 加速内核的性能统计数据和原始算子

请查看 TensorFlow 文档中的 Profiler 指南，详细了解上述工具。

入门指南

快速入门 Profiler 的最佳方法是参照此处的 Colab 教程。我们将在这篇博文中介绍一些重要的步骤和洞察。首先，使用以下命令安装用于 TensorBoard 的 Profiler 插件：

pip install -U tensorboard_plugin_profile

这会为已安装完的 TensorBoard 添加完整的 Profiler 功能。接着，确保模型训练能获取一个分析文件。在本例中，我们将在 Keras 中使用 TensorBoard 回调：

tensorboard_callback = tf.keras.callbacks.TensorBoard (log_dir = logs,
                                                      profile_batch = '500,510')

我们可以使用 profile_batch 参数选择要分析的批次。这样便可选择要捕获的步数（建议不超过 10 个）。同时，这也能帮助我们跳过分析前几个批次，以避免因初始化开销而导致的误差。关于捕获分析文件的其他方法，请参阅此处。接下来，使用以下命令启动 TensorBoard：

tensorboard --logdir {log directory}    # in terminal
%tensorboard --logdir {log directory}   # in Colab

点击**分析文件 **(Profile) 后，我们将看到概述页：

此页面立刻会显示程序的具体性能状况。除了实用的摘要外，此页面还将显示一条建议，告知我们程序存在输入限制（input-bound，说明我们的加速器在等待输入时浪费了时间）。这是一个常见问题。
按照本教程中的说明，我们可将平均步骤时间从约 30 毫秒缩减至约 3 毫秒。换言之，性能提升了 10 倍！虽然这个示例很简单，但 Google 的工程师和研究人员常表示，他们通过某些重要因素实现了性能的显著提升。

建议

性能优化是一个反复迭代的过程，有时可能会令人沮丧，因为您很难精准定位程序中性能瓶颈的位置。我们的新工具 Profiler 不仅能指明程序中存在瓶颈的位置，而且还能提供解决瓶颈的具体建议供参考，以加快代码执行速度。您可以按照提供的建议进行操作，从而缩短优化程序所需的总体时间。
当打开 TensorBoard 查看分析结果时，您可以在概述页的 “步骤” 时间图表下方看到代码优化建议。数据输入流水线的配置不当是代码执行缓慢最常见原因之一。利用 输入流水线分析器 的功能，您可以有效识别和消除数据输入流水线中的瓶颈。请阅读 Profiler 指南的 “最佳实践” 一节，详细了解关于实现最佳性能的可用策略。

TensorFlow Profiler 未来计划

除解决反馈问题之外，我们还在扩展这一分析工具的功能。以下是我们目前正在研究的几个领域：

内存分析器：查看一段时间内的内存使用情况以及相关算子 / 训练步骤。
Keras 分析：实现 Profiler 中的信息与 Keras 的关联。例如，这样便可识别与 Trace Viewer 中显示的算子相对应的具体 Keras 层。
Multiworker GPU 分析：实现对多个 GPU worker 进行分析并汇总分析结果。分析多个 worker 之间的热点和通信情况。

我们很高兴并将持续向社区分享 Google 内部用来提升机器学习性能的工具。如果您急需我们帮忙实现某些特定功能，或者报告错误，请随时在 GitHub 上提出 issue。

原文： Introducing the new TensorFlow Profiler / 中文发布：TensorFlow 公众号

新开源的 TensorFlow 性能分析工具 — TensorFlow Profiler

什么是 TensorFlow Profiler？

入门指南

建议

更多资源

TensorFlow Profiler 未来计划