TensorBoard是检查TensorFlow模型的强大工具，开发者可以通过查看模型结构的预览图确保模型的构建符合其预期。TensorFlow2.0之后将Keras也整合到了框架中，开发者可以在利用Keras搭建模型后利用回调函数来记录模型结构。但是本文将重点介绍如何可视化一个函数的计算过程。

TensorFlow官方文档中指出，可以使用TensorBoard中的TensorFlow Summary Trace API记录签名函数从而进行可视化。
下面是使用Summary Trace API的步骤：

• 使用tf.function定义和注释要进行可视化的函数（tf.function注释后可以将Python计算函数转换为高性能的TensorFlow静态图）
• 在函数调用站点之前立即使用tf.summary.trace_on()
• 通过传递profiler=True将配置文件信息（内存、CPU时间）添加到图中
• 使用摘要文件编写器，调用tf.summary.trace_export()保存日志数据

示例如下所示：

# 定义和注释要进行可视化的函数
@tf.function
def my_func1(x, y):
    return tf.add(x, y)

# 设置时间和日志文件存放路径
time_stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = './desktop/logs/%s' %time_stamp
graph_writer = tf.summary.create_file_writer(logdir)

# 定义数据
num1 = tf.constant(1)
num2 = tf.constant(2)

# 在调用函数前先执行tf.summary.trace_on()函数
tf.summary.trace_on(graph=True, profiler=True)

# 调用函数
result1 = my_func1(num1, num2)

# 保存日志数据
with graph_writer.as_default():
    tf.summary.trace_export(name="my_func_trace", step=0, profiler_outdir=logdir)

最后通过执行如下命令查看可视化结果：

tensorboard --logdir ./desktop/logs