TensorBoard模型图

Posted 2021-08-12 一颗小树x

前言

TensorBoard的图仪表盘（graph）是检测模型的强大工具，可以快速查看模型结构的预览图，并确保其符合预期想法。还可以查看操作级图，以了解Tensorflow程序。

 

一、训练模型并记录数据

首先定义一个简单的顺序模型

# 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'])

 

训练之前，请定义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])

打开TensorBoard

点击顶部的“graph”来选择图形仪表盘。

%tensorboard --logdir logs

 

1.1 op-level 执行图

默认情况下，TensorBoard显示op-level图。（在左侧，可以看到已选择“Default”标签）

请注意，图中的数据是从下到上流动的，因此于代码相比是上到下；

还可以看到该图与Keras模型定义精密匹配，并具有其他计算节点的额外边缘。

通过滚动鼠标来放大或缩小；

单击左键进行拖拉平移；

鼠标移动到节点处，然后双击能看到进一步的结构信息：

 

1.2 概念图

“概念图”是Keras模型的视图。如果需要重新使用保存的模型，并且想要检查或验证其结构，这比较有用。

选择“keras”标签，来查看概念图。

双击结点以查看模型的结构：

 

1.3 完整代码

加载tensorboard

# Load the TensorBoard notebook extension.
%load_ext tensorboard

删除旧日志（可选）

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

 执行如下代码：

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

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."

# 定义一个 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


# 训练模型并记录数据
# 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])

启动 TensorBoard，通过点击顶部的 “graph” 来选择图形仪表板。

%tensorboard --logdir logs

参考：https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/tensorboard/graphs.ipynb#scrollTo=e25E37vd1xEW

 

二、tf.function 可视化图

背景

上面是基于Keras模型的图，其中这些图是用国定义Keras层并调用Model.fit( )创建的。

如果使用tf.function( )函数自定义训练，那如何使用TensorBoard可视化呢？

 

实现

我们可以使用TensorBoard中 TensorFlow Summary Trace API 记录签名函数进行可视化。

先看一下代码：

# 需要跟踪的函数
@tf.function
def my_func(x, y):
  # 一个简单的relu层
  return tf.nn.relu(tf.matmul(x, y))

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

# 函数的示例数据
x = tf.random.uniform((3, 3))
y = tf.random.uniform((3, 3))

# 用 tf.summary.trace_on() 和 tf.summary.trace_export() 来记录函数调用
tf.summary.trace_on(graph=True, profiler=True)

# 跟踪时只调用一个 tf.function
z = my_func(x, y)
with writer.as_default():
  tf.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)

实现流程：

• 使用tf.function 定义训练模型的函数
• 在tf.function函数调用站点前，使用tf.summary.trace_on( )，通过传递profiler = True 将配置文件信息（内存、CPU时间等）添加到图中
• 使用摘要文件编写其，调用tf.summary.trace_export( )保存日志数据
• 打开TensorBoard

 

打开TensorBoard

%tensorboard --logdir logs/func

 

后面的文章后进一步介绍~~