推出 TensorFlow 性能分析工具 — TensorFlow Profiler

文 / 技术撰稿人 Anirudh Sriram 与产品经理 Gal Oshri

性能是机器学习 (ML) 研究和生产解决方案成功与否的关键因素。提升模型训练的速度可加快迭代并减少开销。有时候，这也是实现特定机器学习解决方案的一项基本要求。

但是，要明确需要优化哪些部分却并非易事。我们需要思考是需要优化特定算子 (op)，亦或是输入流水线？

为帮助解答这类问题，我们开发了一套用于 TensorFlow 的综合性能分析工具：TensorFlow Profiler。这套工具除了能捕捉多方面的细节信息辅助分析，还能提供解决性能瓶颈（例如，输入受限程序）的一些建议。

• TensorFlow Profiler
  https://tensorflow.google.cn/tensorboard/tensorboard_profiling_keras

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

TensorFlow Profiler 概述页

   

什么是 TensorFlow Profiler？

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

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

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

• GPU 内核统计数据：显示每个 GPU 加速内核的性能统计数据和原始算子

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

• Profiler 指南
  https://tensorflow.google.cn/guide/profiler#profiler_tools

入门指南

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

pip install -U tensorboard_plugin_profile

• 此处
  https://tensorflow.google.cn/tensorboard/tensorboard_profiling_keras

• 插件
  https://github.com/tensorflow/tensorboard#can-i-make-my-own-plugin

这会为已安装完的 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

• 此处
  https://tensorflow.google.cn/guide/profiler#collect_performance_data

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

此页面立刻会显示程序的具体性能状况。除了实用的摘要外，此页面还将显示一条建议，告知我们程序存在输入限制（input-bound，说明我们的加速器在等待输入时浪费了时间）。这是一个常见问题。

按照本教程中的说明，我们可将平均步骤时间从约 30 毫秒缩减至约 3 毫秒。换言之，性能提升了 10 倍！虽然这个示例很简单，但 Google 的工程师和研究人员常表示，他们通过某些重要因素实现了性能的显著提升。

建议

性能优化是一个反复迭代的过程，有时可能会令人沮丧，因为您很难精准定位程序中性能瓶颈的位置。我们的新工具 Profiler 不仅能指明程序中存在瓶颈的位置，而且还能提供解决瓶颈的具体建议供参考，以加快代码执行速度。您可以按照提供的建议进行操作，从而缩短优化程序所需的总体时间。

当打开 TensorBoard 查看分析结果时，您可以在概述页的“步骤”时间图表下方看到代码优化建议。数据输入流水线的配置不当是代码执行缓慢最常见原因之一。利用 输入流水线分析器 的功能，您可以有效识别和消除数据输入流水线中的瓶颈。请阅读 Profiler 指南的“最佳实践”一节，详细了解关于实现最佳性能的可用策略。

• 最佳实践
  https://tensorflow.google.cn/guide/profiler#best_practices_for_optimal_model_performance

更多资源

详情请查看以下资源：

• Colab 中的 Profiler 教程：

  https://tensorflow.google.cn/tensorboard/tensorboard_profiling_keras

• 详细指南：

  https://tensorflow.google.cn/guide/profiler

• GitHub 代码库

  https://github.com/tensorflow/profiler

• TensorFlow Dev Summit 2020 演讲：

  https://space.bilibili.com/64169458/channel/detail?cid=41910

中文字幕 | TensorFlow 2.x 版本的性能检测

TensorFlow Profiler 未来计划

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

• 内存分析器：查看一段时间内的内存使用情况以及相关算子/训练步骤。
• Keras 分析：实现 Profiler 中的信息与 Keras 的关联。例如，这样便可识别与 Trace Viewer 中显示的算子相对应的具体 Keras 层。

• Multiworker GPU 分析：实现对多个 GPU worker 进行分析并汇总分析结果。分析多个 worker 之间的热点和通信情况。

我们很高兴并将持续向社区分享 Google 内部用来提升机器学习性能的工具。如果您急需我们帮忙实现某些特定功能，或者报告错误，请随时在 Github (https://github.com/tensorflow/profiler/issues) 提 issue。

— 推荐阅读 —

请点击 “ 阅读原文 ”，观看更多 TensorFlow 视频 。