检查 TensorFlow 图

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 Github 上查看源代码

下载笔记本

概述

TensorBoard 的 Graphs 信息中心是检查 TensorFlow 模型的强大工具。您可以快速查看模型结构的概念计算图，并确保其符合您的预期设计。您还可以查看运算级计算图以了解 TensorFlow 如何理解您的程序。检查运算级计算图可以使您深入了解如何更改模型。例如，如果训练进度比预期的慢，则可以重新设计模型。

本教程简要概述了如何生成计算图诊断数据并将其在 TensorBoard 的 Graphs 信息中心中呈现。您将为 Fashion-MNIST 数据集定义和训练一个简单的 Keras 序贯模型，并学习如何记录和检查模型计算图。您还将使用跟踪 API 为使用新 tf.function 注解创建的函数生成计算图数据。

设置

# Load the TensorBoard notebook extension.
%load_ext tensorboard

from datetime import datetime
from packaging import version

import tensorflow as tf
from tensorflow import keras

print("TensorFlow version: ", tf.__version__)
assert version.parse(tf.__version__).release[0] >= 2, \
    "This notebook requires TensorFlow 2.0 or above."

TensorFlow version:  2.2.0

import tensorboard
tensorboard.__version__

'2.2.1'

# Clear any logs from previous runs
rm -rf ./logs/

定义一个 Keras 模型

在此示例中，分类器是一个简单的四层顺序模型。

# Define the model.
model = keras.models.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(32, activation='relu'),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10, activation='softmax')
])

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'])

下载并准备训练数据

(train_images, train_labels), _ = keras.datasets.fashion_mnist.load_data()
train_images = train_images / 255.0

训练模型并记录数据

训练之前，请定义 Keras TensorBoard callback, 并指定日志目录。通过将此回调传递给 Model.fit(), 可以确保在 TensorBoard 中记录图形数据以进行可视化。

# Define the Keras TensorBoard callback.
logdir="logs/fit/" + datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)

# Train the model.
model.fit(
    train_images,
    train_labels, 
    batch_size=64,
    epochs=5, 
    callbacks=[tensorboard_callback])

Epoch 1/5
938/938 [==============================] - 2s 2ms/step - loss: 0.6955 - accuracy: 0.7618
Epoch 2/5
938/938 [==============================] - 2s 2ms/step - loss: 0.4877 - accuracy: 0.8296
Epoch 3/5
938/938 [==============================] - 2s 2ms/step - loss: 0.4458 - accuracy: 0.8414
Epoch 4/5
938/938 [==============================] - 2s 2ms/step - loss: 0.4246 - accuracy: 0.8476
Epoch 5/5
938/938 [==============================] - 2s 2ms/step - loss: 0.4117 - accuracy: 0.8508
<tensorflow.python.keras.callbacks.History at 0x7f656ecc3fd0>

op-level graph

启动 TensorBoard，然后等待几秒钟以加载 UI。通过点击顶部的 “graph” 来选择图形仪表板。

%tensorboard --logdir logs

您还可以选择使用 TensorBoard.dev 来创建托管的可共享实验。

!tensorboard dev upload \
  --logdir logs \
  --name "Sample op-level graph" \
  --one_shot

默认情况下，TensorBoard 显示 op-level图。（在左侧，您可以看到已选择 “Default” 标签。）请注意，图是倒置的。数据从下到上流动，因此与代码相比是上下颠倒的。但是，您可以看到该图与 Keras 模型定义紧密匹配，并具有其他计算节点的额外边缘。

图通常很大，因此您可以操纵图的可视化效果：

滚动到 zoom 来放大和缩小
拖到 pan 平移
双击切换 node expansion 来进行节点扩展（一个节点可以是其他节点的容器）

您还可以通过单击节点来查看元数据。这使您可以查看输入，输出，形状和其他详细信息。

概念图

除了执行图，TensorBoard 还显示一个“概念图”。这只是 Keras 模型的视图。如果您要重新使用保存的模型并且想要检查或验证其结构，这可能会很有用。

要查看概念计算图，请选择“keras”标签。在此示例中，您将看到一个收起的 Sequential 节点。双击该节点可以查看模型的结构：

tf.functions 的图

目前为止的示例描述了 Keras 模型的计算图，其中这些计算图是通过定义 Keras 层并调用 Model.fit() 创建的。

您可能会遇到需要使用 tf.function 注释来[autograph]的情况，即将 Python 计算函数转换为高性能 TensorFlow 图。对于这些情况，您可以使用 TensorBoard 中的 TensorFlow Summary Trace API 记录签名函数以进行可视化。

要使用 Summary Trace API ，请执行以下操作:

使用 tf.function 定义和注释功能
在函数调用站点之前立即使用 tf.summary.trace_on()
通过传递 profiler=True 将性能剖析信息（内存、CPU 时间）添加到计算图中
使用摘要文件编写器，调用 tf.summary.trace_export() 保存日志数据

然后，您可以使用 TensorBoard 查看函数的行为。

# The function to be traced.
@tf.function
def my_func(x, y):
  # A simple hand-rolled layer.
  return tf.nn.relu(tf.matmul(x, y))

# Set up logging.
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = 'logs/func/%s' % stamp
writer = tf.summary.create_file_writer(logdir)

# Sample data for your function.
x = tf.random.uniform((3, 3))
y = tf.random.uniform((3, 3))

# Bracket the function call with
# tf.summary.trace_on() and tf.summary.trace_export().
tf.summary.trace_on(graph=True, profiler=True)
# Call only one tf.function when tracing.
z = my_func(x, y)
with writer.as_default():
  tf.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)

%tensorboard --logdir logs/func

现在，您可以看到 TensorBoard 所理解的函数结构。单击 “Profile” 按钮以查看 CPU 和内存统计信息。