跳到主要内容

在 TensorBoard 之中查看模型结构图

在之前的学习过程之中，我们学习了如何自定义查看训练过程之中的各项指标。在实际的应用过程之中，为了保证模型构建的准确性，我们也会经常查看网络的模型结构图。那么这节课我们就来看一下如何在 TensorBoard 之中查看模型图。

1. 如何在 TensorBoard 之中生成 Keras 模型结构图​

倘若我们通过 tf.keras API 来自定义了一个网络模型，那么我们在 TensorBoard 来查看模型图是非常简单的一件事情。

当我们使用 tf.keras 的模型的 fit() 方法的时候，框架会自动帮我们绘制模型结构图。

如下代码所示：

首先我们定义模型、数据与相应的参数。

import tensorflow as tf

(x_train, y_train),(x_test, y_test) = tf.keras.datasets.mnist.load_data()

x_train, x_test = x_train / 255.0, x_test / 255.0

def create\_model():
  return tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(256, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
  ])

model = create_model()
model.compile(optimizer='adam',
              loss='sparse\_categorical\_crossentropy',
              metrics=[])

然后我们定义相应的 TensorBoard 日志目录，同时对模型使用 fit() 进行训练：

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs')

model.fit(x=x_train, y=y_train, 
          epochs=3, 
          validation_data=(x_test, y_test), 
          callbacks=[tensorboard_callback])

最后我们就可以打开 TensorBoard 并在浏览器查看：

tensorboard --logdir logs

我们就可以在浏览器的 Graph 标签页之中看到模型图了：

图片描述

2. 如何在 TensorBoard 之中生成使用 tf.function 函数定义的图​

在实际的应用过程之中，有很多的情况下，我们需要使用 tf.function 来加速模型的速度并自定义训练过程。那么这个时候我们要如何才能查看网络的模型结构图呢？

其实也很简单，我们只需经过如下几个步骤：

  • 确保 tf.function 函数修饰了我们需要进行可视化的操作，这边就是模型的过程；
  • 创建一个 TensorBoard 的日志写入器 tf.summary.create_file_writer() ；
  • 通过 tf.summary.trace_on() API 进行变量路径的追踪；
  • 执行我们需要可视化的操作；
  • 使用 tf.summary.trace_export() API 将图写入日志。

在这里，我们可以使用一个很简单的例子来查看操作的结构：


# 定义网络的操作
@tf.function
def test\_func(x, y):
  z = tf.matmul(x, y)
  z = z \* 5.0
  z = tf.nn.relu(z)
  return z

# 创建写入器
writer = tf.summary.create_file_writer('./logs/3')

# 创建初试数据
x = tf.random.uniform((5, 5))
y = tf.random.uniform((5, 5))

# 开启变量追踪
tf.summary.trace_on(graph=True, profiler=True)

# 运行程序
z = test_func(x, y)

# 将日志输出
with writer.as_default():
  tf.summary.trace_export(
      name="test\_func\_graph",
      step=1,
      profiler_outdir='./logs/3')

在这里，我们首先定义了一个基本的模型操作，该模型操作由一个矩阵乘法、一个常量乘法、外加一个 Relu 激活层组成。

在运行完操作之后，我们便使用 tf.summary.trace_export() API 来将模型图输入道日志之中。

然后我们便可以在浏览器之中查看到相应的模型图：

图片描述

可以看到，该模型图完整的反映了我们的操作。

3. TensorBoard 之中基本、基本的操作​

既然了解了如何将模型图输出到日志，那么接下来我们就应该查看在 TensorBoard 之中对模型图的基本操作。

3.1 平移、缩放以及详细信息的查看​

在 TensorBoard 之中，使用鼠标滚轮即可实现模型图的缩放，当我们一直放大，会看到操作内部的细节。

并且按住鼠标左键，移动鼠标，即可实现模型图的移动操作。

双击网络节点，即可展开网络节点，从而查看到网络内部的细节操作。

3.2 模型的节点的搜索​

图片描述

在左侧的最上方，可以搜索自己想要查看的节点，这里是支持正则表达式的。

3.3 模型的下载​

点击左侧的 Download PNG 即可下载带有透明度的、网络模型的图片。

3.4 切换网络模型​

点击左侧的 Run 按钮，即可选择不同的网络模型进行查看，前提是我们已经将网络模型输入到日志之中去。

3.5 切换查看方式​

点击左侧的 Tag 选项，即可查看网络的查看方式。

默认是查看中粒度的网络模型，如果我们的模型是使用 Keras 定义的，那么我们可以选择查看 Keras 结构，这是一个总体的概览，可以帮助我们掌握大体的网络结构。

3.6 模型图的图例​

当我们遇到一些不理解的图标的时候，我们可以通过左下角的图例进行查询：

图片描述

4. 小结​

在这节课之中，我们学习率如何在 TensorBoard 之中查看 Keras 模型，同时也了解了如何产看自定义的操作过程，最后我们了解了 TensorBoard 的一些基本操作。 TensorBoard 也在持续更新，未来一定会有更多新的功能。