怎么在Python中可视化TensorFlow模型的计算图结构

用 tf.summary.trace_export 导出 GraphDef 供 TensorBoard 查看

TensorFlow 2.x 默认启用 Eager Execution，计算图不显式构建，直接调用 tf.summary.trace_on 和 tf.summary.trace_export 是最轻量、兼容性最好的方式。它不依赖模型是否继承 tf.keras.Model，只要在 eager 模式下执行过一次前向传播即可。

• 必须在 tf.summary.trace_on 后立即运行至少一次模型调用（如 model(x)），否则导出为空
• 导出路径需为完整目录（如 "./logs/trace"），TensorBoard 会自动识别其中的 plugins/profile/ 子结构
• 若用 tf.function 包装模型，需确保 trace 在 @tf.function 内部或外层统一开启，否则可能只捕获 Python 层调用
• 导出后启动 TensorBoard：tensorboard --logdir=./logs/trace --bind_all，然后访问 http://localhost:6006/#graphs

对 Keras 模型使用 tf.keras.utils.plot_model 生成静态结构图

这个函数画的是模型的层连接拓扑（Layer Graph），不是运行时的实际计算图（GraphDef），适合快速确认输入输出形状、层类型和连接关系，但不反映 control flow（如 tf.cond）、梯度计算路径或变量绑定细节。

• show_shapes=True 能显示每层输入/输出 shape，但对动态 batch size（如 None）只显示占位符
• expand_nested=True 可展开 tf.keras.Sequential 或子模型，但对自定义 tf.keras.layers.Layer 中的内部 tf.function 无效
• 输出格式依赖 pydot 和 graphviz：需系统安装 graphviz，Python 安装 pydot，否则报错 ImportError: Failed to import pydot
• 生成的 .png 文件不包含 op-level 信息（如 AddV2、MatMul），仅含 Keras 层名（如 Conv2D、Dense）

手动提取并可视化 ConcreteFunction 的底层 GraphDef

当需要 inspect 实际编译后的计算图（比如调试 tf.function 内部优化、查看 fused ops 或 control dependencies），必须从 ConcreteFunction 提取 graph.as_graph_def()，再用第三方工具解析。

• 获取方式：cf = model.call.get_concrete_function(input_spec); graph_def = cf.graph.as_graph_def()，注意 get_concrete_function 必须传入与实际输入匹配的 tf.TensorSpec
• graph_def 是 protobuf 格式，不能直接 human-read；可用 tf.train.write_graph(graph_def, './graph', 'model.pbtxt', as_text=True) 输出可读文本
• 文本文件中大量出现 _XlaLaunch、StatefulPartitionedCall 等封装 op，说明 XLA 或函数内联已生效，此时原始层名可能被抹除
• 想图形化查看？可转成 ONNX（用 tf2onnx）再用 Netron 打开，但部分 TF 特有 op（如 ResourceGather）可能丢失或报错

避免常见陷阱：Eager vs Graph 模式混淆导致图为空

很多用户发现 TensorBoard 里 graphs 页面空白，或 plot_model 报 AttributeError: 'Tensor' object has no attribute 'shape'，根本原因在于混用了 eager 执行和图构建预期。

• 不要在 tf.function 外部调用 tf.summary.trace_on 却在内部执行模型——trace 必须覆盖实际执行路径
• Keras 模型若用 model.compile(..., run_eagerly=True)，plot_model 正常，但 trace_export 会失效（因为没触发图构建）
• 检查当前模式：tf.executing_eagerly() 返回 True 表示 eager 模式；若为 False，说明意外进入了 legacy graph mode，需排查是否调用了 tf.compat.v1.disable_eager_execution()
• 对 SavedModel 加载的模型，tf.saved_model.load 返回的是签名函数集合，需先获取具体 signature（如 infer = loaded.signatures['serving_default']），再对其调用 get_concrete_function

真正要看到训练时反向传播路径、variable update op 顺序、gradient accumulation 结构，必须用 tf.summary.trace_export 配合 profiler 插件，并在训练 step 中开启 trace；静态图工具永远看不到这些动态行为。