TensorFlow Eager模式下能否提取底层图用于TensorBoard可视化？

好问题！确实，Eager模式的设计就是让你摆脱手动构建Graph和Session的繁琐，但底层确实会动态生成临时Graph用于执行——不过这些Graph默认是用完就销毁的，没法直接提取现成的Session或者持久化Graph。但完全不需要把整个模型重构为传统的Graph/Session模式，下面给你两种实用的方法：

方法1：用tf.function包装代码生成可导出的Graph #

tf.function是TensorFlow用来将Eager风格的代码自动转换为Graph模式（AutoGraph）的工具，你只需要把要可视化的模型逻辑用它包裹，就能生成稳定的Graph，然后导出给TensorBoard：

import tensorflow as tf

# 定义一个Eager模式下的简单模型
class MyModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(10, activation='relu')
        self.dense2 = tf.keras.layers.Dense(2)
    
    # 用tf.function包装前向传播逻辑
    @tf.function
    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

# 初始化模型和输入
model = MyModel()
sample_input = tf.random.normal((1, 32))

# 开启Trace，捕获Graph信息
tf.summary.trace_on(graph=True, profiler=True)
# 执行一次模型触发Graph构建
model(sample_input)
# 导出Trace到日志目录
with tf.summary.create_file_writer('./logs/eager_graph').as_default():
    tf.summary.trace_export(
        name="my_eager_model_trace",
        step=0,
        profiler_outdir='./logs/eager_graph'
    )

运行完这段代码后，你只需要启动TensorBoard（tensorboard --logdir=./logs/eager_graph），就能在"Graphs"标签页看到模型的可视化图了。

方法2：直接捕获Eager执行过程的临时Graph #

如果不想用tf.function包装，你也可以直接在Eager执行时开启Trace，捕获单次执行的操作图：

import tensorflow as tf

# 同样定义模型（不需要tf.function）
class MyModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(10, activation='relu')
        self.dense2 = tf.keras.layers.Dense(2)
    
    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

model = MyModel()
sample_input = tf.random.normal((1, 32))

# 开启Trace，指定捕获Graph
tf.summary.trace_on(graph=True)
# 执行模型前向传播
model(sample_input)
# 导出Trace
with tf.summary.create_file_writer('./logs/eager_instant_graph').as_default():
    tf.summary.trace_export(name="instant_eager_trace", step=0)

这种方法会捕获你这次执行过程中所有的TensorFlow操作，但注意它是单次执行的临时图，可能包含一些Eager模式下的临时操作，不如tf.function生成的Graph简洁规范。

注意事项 #

• 不管用哪种方法，都需要至少执行一次模型（或者被tf.function包装的函数），因为Graph是在首次执行时才会被构建出来的。
• 如果你的模型有多个分支或者动态控制流（比如if/else循环），tf.function会根据输入的形状自动生成对应的Graph，你可以用不同的输入多次执行来捕获完整的Graph逻辑。

总之，不需要把整个模型重构为传统的Graph/Session模式，用上面的方法就能轻松导出可用于TensorBoard可视化的Graph啦！

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