TensorFlow 2.0中无需计算即可在TensorBoard可视化模型图的方法

既然你已经通过tf.compat.v1.disable_eager_execution()切换到了图模式，其实只需要添加几步就能把模型图写入TensorBoard，完全不需要执行任何计算。下面给你两种可行的方案，适配不同的使用习惯：

方案1：使用TF2原生的Summary API #

这是更贴合TF2风格的做法，通过跟踪图结构并导出的方式实现：

from nets import i3d
import tensorflow as tf

def i3d_output(model, x):
    out, _ = model(x)
    return out

# 禁用eager execution，切换到图模式
tf.compat.v1.disable_eager_execution()

# 开启图跟踪，只记录图结构（不需要性能分析的话profiler设为False）
tf.summary.trace_on(graph=True, profiler=False)

# 构建你的输入和模型（这一步会生成计算图）
x = tf.random.uniform(shape=(4,179,224,224,3))
model = i3d.InceptionI3d()
net = i3d_output(model, x)

# 创建Summary Writer
train_summary_writer = tf.summary.create_file_writer('/home/uujjwal/bmvc2019')

# 将跟踪到的图结构写入TensorBoard
with train_summary_writer.as_default():
    tf.summary.trace_export(
        name='i3d_model_graph',  # 给你的图起个名字
        step=0,  # 不需要训练步骤，设为0即可
        profiler_outdir='/home/uujjwal/bmvc2019'
    )

关键步骤解释： #

• tf.summary.trace_on(graph=True)：告诉TensorFlow开始记录当前上下文的计算图结构。
• tf.summary.trace_export：把记录好的图结构导出并写入到指定的Writer中，完成后会自动停止跟踪。

方案2：兼容TF1.X的写法（更接近你熟悉的flush方式） #

如果你习惯TF1.X的操作逻辑，可以用兼容API直接把当前图添加到Writer：

from nets import i3d
import tensorflow as tf

def i3d_output(model, x):
    out, _ = model(x)
    return out

# 禁用eager execution，切换到图模式
tf.compat.v1.disable_eager_execution()

# 构建你的输入和模型
x = tf.random.uniform(shape=(4,179,224,224,3))
model = i3d.InceptionI3d()
net = i3d_output(model, x)

# 创建TF1风格的FileWriter并添加当前图
writer = tf.compat.v1.summary.FileWriter('/home/uujjwal/bmvc2019')
writer.add_graph(tf.compat.v1.get_default_graph())
# 手动刷新并关闭Writer，确保图结构写入文件
writer.flush()
writer.close()

这种方式和你在TF1.X里的操作逻辑几乎一致，因为禁用eager后，tf.compat.v1.get_default_graph()可以直接拿到当前生成的计算图，添加到Writer后刷新即可。

完成后，你只需要在终端运行TensorBoard命令：

tensorboard --logdir=/home/uujjwal/bmvc2019

就能在浏览器里看到模型的图结构了。

内容的提问来源于stack exchange，提问作者Ujjwal