【tensorboard】可视化模型

关键使用tf.summary.trace_on(graph=True, profiler=True)跟踪张量的流动，tf.summary.trace_export导出图的结构。

import tensorflow as tf
import os
# 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.
logdir = 'logs'
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)
if not os.path.exists('{}/plugins/profile'.format(logdir)):
    os.mkdir('{}/plugins/'.format(logdir))
    os.mkdir('{}/plugins/profile'.format(logdir))
with writer.as_default():
    tf.summary.trace_export(
        name="my_func_trace",
        step=0,
        profiler_outdir=logdir) # 导出图结构

命令行输入tensorboard --logdir logs登录tensorboard，可以在Graphs中看到网络的结构