TensorFlow Profiler:分析模型性能

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

概述

机器学习算法通常计算量很大。因此,量化机器学习应用程序的性能至关重要,以确保您正在运行模型的最佳优化版本。使用 TensorFlow Profiler 分析 TensorFlow 代码的执行情况。

设置

from datetime import datetime
from packaging import version

import os

TensorFlow Profiler 需要最新版本的 TensorFlow 和 TensorBoard (>=2.2)。

pip install -U tensorboard_plugin_profile
import tensorflow as tf

print("TensorFlow version: ", tf.__version__)
TensorFlow version:  2.2.0-dev20200405

确认 TensorFlow 可以访问 GPU。

device_name = tf.test.gpu_device_name()
if not device_name:
  raise SystemError('GPU device not found')
print('Found GPU at: {}'.format(device_name))
Found GPU at: /device:GPU:0

使用 TensorBoard 回调训练图像分类模型

在本教程中,您将通过捕获训练模型以对 MNIST 数据集 中的图像进行分类而获得的性能配置文件,来探索 TensorFlow Profiler 的功能。

使用 TensorFlow 数据集导入训练数据并将其拆分为训练集和测试集。

import tensorflow_datasets as tfds
tfds.disable_progress_bar()
(ds_train, ds_test), ds_info = tfds.load(
    'mnist',
    split=['train', 'test'],
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
)
WARNING:absl:Dataset mnist is hosted on GCS. It will automatically be downloaded to your
local data directory. If you'd instead prefer to read directly from our public
GCS bucket (recommended if you're running on GCP), you can instead set
data_dir=gs://tfds-data/datasets.
Downloading and preparing dataset mnist/3.0.0 (download: 11.06 MiB, generated: Unknown size, total: 11.06 MiB) to /root/tensorflow_datasets/mnist/3.0.0...
Dataset mnist downloaded and prepared to /root/tensorflow_datasets/mnist/3.0.0. Subsequent calls will reuse this data.

通过将像素值归一化为 0 到 1 之间的值来预处理训练数据和测试数据。

def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label

ds_train = ds_train.map(normalize_img)
ds_train = ds_train.batch(128)
ds_test = ds_test.map(normalize_img)
ds_test = ds_test.batch(128)

使用 Keras 创建图像分类模型。

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
  tf.keras.layers.Dense(128,activation='relu'),
  tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer=tf.keras.optimizers.Adam(0.001),
    metrics=['accuracy']
)

创建一个 TensorBoard 回调以捕获性能配置文件,并在训练模型时调用它。

# Create a TensorBoard callback
logs = "logs/" + datetime.now().strftime("%Y%m%d-%H%M%S")

tboard_callback = tf.keras.callbacks.TensorBoard(log_dir = logs,
                                                 histogram_freq = 1,
                                                 profile_batch = '500,520')

model.fit(ds_train,
          epochs=2,
          validation_data=ds_test,
          callbacks = [tboard_callback])
Epoch 1/2
469/469 [==============================] - 11s 22ms/step - loss: 0.3684 - accuracy: 0.8981 - val_loss: 0.1971 - val_accuracy: 0.9436
Epoch 2/2
 50/469 [==>...........................] - ETA: 9s - loss: 0.2014 - accuracy: 0.9439WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/summary_ops_v2.py:1271: stop (from tensorflow.python.eager.profiler) is deprecated and will be removed after 2020-07-01.
Instructions for updating:
use `tf.profiler.experimental.stop` instead.
WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/summary_ops_v2.py:1271: stop (from tensorflow.python.eager.profiler) is deprecated and will be removed after 2020-07-01.
Instructions for updating:
use `tf.profiler.experimental.stop` instead.
469/469 [==============================] - 11s 24ms/step - loss: 0.1685 - accuracy: 0.9525 - val_loss: 0.1376 - val_accuracy: 0.9595
<tensorflow.python.keras.callbacks.History at 0x7f23919a6a58>

使用 TensorFlow Profiler 分析模型训练性能

TensorFlow Profiler 嵌入在 TensorBoard 中。使用 Colab 魔法加载 TensorBoard 并启动它。通过导航到 **分析** 选项卡来查看性能配置文件。

# Load the TensorBoard notebook extension.
%load_ext tensorboard

此模型的性能配置文件类似于下面的图像。

# Launch TensorBoard and navigate to the Profile tab to view performance profile
%tensorboard --logdir=logs
<IPython.core.display.Javascript object>

**分析** 选项卡将打开概述页面,该页面显示模型性能的高级摘要。查看右侧的步长时间图,您可以看到该模型高度依赖输入(即,它在数据输入管道中花费了大量时间)。概述页面还为您提供有关优化模型性能的潜在下一步操作的建议。

要了解输入管道中性能瓶颈发生的位置,请从左侧的 **工具** 下拉菜单中选择 **跟踪查看器**。跟踪查看器显示了在分析期间 CPU 和 GPU 上发生的各种事件的时间线。

跟踪查看器在垂直轴上显示多个事件组。每个事件组都有多个水平轨道,其中填充了跟踪事件。轨道是线程或 GPU 流上执行的事件的事件时间线。单个事件是时间线轨道上的彩色矩形块。时间从左到右移动。使用键盘快捷键 W(放大)、S(缩小)、A(向左滚动)和 D(向右滚动)来导航跟踪事件。

单个矩形表示跟踪事件。选择浮动工具栏中的鼠标光标图标(或使用键盘快捷键 1),然后单击跟踪事件以分析它。这将显示有关事件的信息,例如其开始时间和持续时间。

除了单击之外,您还可以拖动鼠标以选择一组跟踪事件。这将为您提供该区域中所有事件的列表以及事件摘要。使用 M 键测量所选事件的时间持续时间。

跟踪事件从以下位置收集

  • CPU: CPU 事件显示在名为 /host:CPU 的事件组下。每个轨道代表 CPU 上的一个线程。CPU 事件包括输入管道事件、GPU 操作 (op) 调度事件、CPU op 执行事件等。
  • GPU: GPU 事件显示在以 /device:GPU: 为前缀的事件组下。每个事件组代表 GPU 上的一个流。

调试性能瓶颈

使用跟踪查看器来定位输入管道中的性能瓶颈。下面的图像显示了性能配置文件的快照。

profiler_trace_viewer_bad_ip

查看事件跟踪,您可以看到 GPU 在 tf_data_iterator_get_next op 在 CPU 上运行时处于非活动状态。此 op 负责处理输入数据并将其发送到 GPU 进行训练。一般来说,最好始终让设备(GPU/TPU)保持活动状态。

使用 tf.data API 优化输入管道。在本例中,让我们缓存训练数据集并预取数据,以确保 GPU 始终有可供处理的数据。有关使用 tf.data 优化输入管道的更多详细信息,请参见 此处

(ds_train, ds_test), ds_info = tfds.load(
    'mnist',
    split=['train', 'test'],
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
)
ds_train = ds_train.map(normalize_img)
ds_train = ds_train.batch(128)
ds_train = ds_train.cache()
ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)
ds_test = ds_test.map(normalize_img)
ds_test = ds_test.batch(128)
ds_test = ds_test.cache()
ds_test = ds_test.prefetch(tf.data.experimental.AUTOTUNE)

再次训练模型,并通过重复使用之前的回调来捕获性能配置文件。

model.fit(ds_train,
          epochs=2,
          validation_data=ds_test,
          callbacks = [tboard_callback])
Epoch 1/2
469/469 [==============================] - 10s 22ms/step - loss: 0.1194 - accuracy: 0.9658 - val_loss: 0.1116 - val_accuracy: 0.9680
Epoch 2/2
469/469 [==============================] - 1s 3ms/step - loss: 0.0918 - accuracy: 0.9740 - val_loss: 0.0979 - val_accuracy: 0.9712
<tensorflow.python.keras.callbacks.History at 0x7f23908762b0>

重新启动 TensorBoard 并打开 **分析** 选项卡以观察更新后的输入管道的性能配置文件。

具有优化输入管道的模型的性能配置文件类似于下面的图像。

%tensorboard --logdir=logs
Reusing TensorBoard on port 6006 (pid 750), started 0:00:12 ago. (Use '!kill 750' to kill it.)
<IPython.core.display.Javascript object>

从概述页面,您可以看到平均步长时间和输入步长时间都减少了。步长时间图还表明该模型不再高度依赖输入。打开跟踪查看器以检查具有优化输入管道的跟踪事件。

profiler_trace_viewer_good_ip

跟踪查看器显示 tf_data_iterator_get_next op 的执行速度快得多。因此,GPU 获得了稳定流的数据来执行训练,并通过模型训练实现了更好的利用率。

总结

使用 TensorFlow Profiler 分析和调试模型训练性能。阅读 Profiler 指南 并观看 TensorFlow Dev Summit 2020 上的 TF 2 中的性能分析 演讲,以了解有关 TensorFlow Profiler 的更多信息。