如何将自定义非Keras TensorFlow模型的变量摘要转为可视化图？

你手上有一个结构未知的自定义TensorFlow非Keras模型，目前只能通过slim.model_analyzer拿到变量摘要，想要生成类似graph.pbtxt对应的层级结构可视化PNG，以下是几种可行方案：

方案1：利用TensorBoard直接生成计算图可视化（推荐） #

如果能在模型运行时捕获计算图，这是最接近graph.pbtxt可视化效果的方法：

1. 在模型代码中添加日志写入逻辑：
   在模型构建完成或会话启动后，添加以下代码将计算图写入日志目录：

   # 假设你已经构建了模型的计算图
   log_dir = "./tf_logs"
   writer = tf.summary.FileWriter(log_dir, tf.get_default_graph())
   writer.close()
   

2. 启动TensorBoard并导出PNG：
   在终端执行命令启动TensorBoard：

   tensorboard --logdir=./tf_logs
   

   打开浏览器访问提示的地址，进入Graphs标签页，找到你的模型图后，点击右上角的下载按钮即可导出PNG格式的可视化图。

方案2：基于变量摘要手动生成层级结构可视化 #

如果无法捕获完整计算图，可从变量名称的层级路径（如model/temb/dense0/W:0）推断模型结构，用graphviz库绘制：

1. 安装依赖库：

   pip install graphviz
   

2. 编写绘图代码：
   解析变量名称的层级关系，生成树形结构：

   import graphviz
   
   # 你的变量摘要文本（可直接读取输出或从文件导入）
   var_summary = """
   model/temb/dense0/W:0 (float32_ref 128x512) [65536, bytes: 262144]
   model/temb/dense0/b:0 (float32_ref 512) [512, bytes: 2048]
   model/temb/dense1/W:0 (float32_ref 512x512) [262144, bytes: 1048576]
   model/temb/dense1/b:0 (float32_ref 512) [512, bytes: 2048]
   model/conv_in/W:0 (float32_ref 3x3x1x128) [1152, bytes: 4608]
   model/conv_in/b:0 (float32_ref 128) [128, bytes: 512]
   model/down_0/block_0/norm1/beta:0 (float32_ref 128) [128, bytes: 512]
   model/down_0/block_0/norm1/gamma:0 (float32_ref 128) [128, bytes: 512]
   model/down_0/block_0/conv1/W:0 (float32_ref 3x3x128x128) [147456, bytes: 589824]
   model/down_0/block_0/conv1/b:0 (float32_ref 128) [128, bytes: 512]
   """
   
   # 解析变量路径
   nodes = set()
   edges = []
   for line in var_summary.strip().split('\n'):
       if not line:
           continue
       var_name = line.split(' ')[0]
       # 拆分层级路径（去掉末尾的:0）
       path_parts = var_name[:-2].split('/')
       # 构建父节点关系
       for i in range(1, len(path_parts)):
           parent = '/'.join(path_parts[:i])
           child = '/'.join(path_parts[:i+1])
           nodes.add(parent)
           nodes.add(child)
           edges.append((parent, child))
   
   # 生成可视化图
   dot = graphviz.Digraph(comment='Model Structure')
   for node in nodes:
       dot.node(node)
   for parent, child in edges:
       dot.edge(parent, child)
   
   # 保存为PNG
   dot.render('model_structure', format='png', cleanup=True)
   

   执行后会生成model_structure.png，展示模型的层级模块结构。

方案3：从检查点恢复计算图并可视化 #

如果有模型的检查点文件（.ckpt），可恢复变量并重建计算图后可视化：

1. 读取检查点并重建计算图：

   import tensorflow as tf
   
   ckpt_path = "./your_model.ckpt"
   # 读取变量信息
   reader = tf.train.NewCheckpointReader(ckpt_path)
   var_to_shape_map = reader.get_variable_to_shape_map()
   
   # 手动构建简化计算图（根据变量名称推断层类型）
   graph = tf.Graph()
   with graph.as_default():
       for var_name, shape in var_to_shape_map.items():
           tf.get_variable(var_name, shape=shape)
       # 保存计算图
       writer = tf.summary.FileWriter("./ckpt_logs", graph)
       writer.close()
   

2. 用TensorBoard导出PNG：
   同方案1，启动TensorBoard后在Graphs标签页下载PNG。

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