如何在TensorFlow 1.x图中移除输入解码节点并对接ResNet？

优先选择的文件格式：frozen graph.pb #

• .ckpt仅保存变量值，没有完整的计算图拓扑结构，无法直接修改节点连接
• saved_model.pb包含图结构和变量，但修改拓扑需要重新构建并导出，流程繁琐
• frozen graph.pb是将所有变量转为常量的完整计算图，结构固定，便于直接编辑节点连接，是最适合的格式。如果只有.ckpt文件，需要先将其转换为冻结图再操作。

具体操作步骤 #

1. 加载原冻结图并分析节点结构 #

首先需要明确原计算图中解码模块的节点、ResNet的输入节点名称。可以通过代码遍历节点或TensorBoard可视化查看：

import tensorflow as tf

# 加载原冻结图
with tf.gfile.GFile('original_frozen_graph.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

# 将图导入默认计算图
tf.import_graph_def(graph_def, name='')

# 遍历所有节点，打印名称以定位关键节点
for op in tf.get_default_graph().get_operations():
    print(op.name)

执行后，你需要找到：解码模块的输出节点（比如decode/image_output）、ResNet的输入节点（比如resnet/input_tensor）。

2. 编辑计算图，替换输入并删除解码模块 #

使用TensorFlow 1.x的tf.contrib.graph_editor工具直接修改节点连接，移除不必要的模块：

import tensorflow.contrib.graph_editor as ge

# 定义新的输入节点，尺寸和数据类型要和ResNet原输入完全匹配
new_input = tf.placeholder(tf.float32, shape=[None, 224, 224, 3], name='custom_input')

# 获取ResNet的输入操作
resnet_input_op = tf.get_default_graph().get_operation_by_name('resnet/input_tensor')
# 找到原输入边（来自解码模块的输出）
original_input_edge = ge.sgv(resnet_input_op).inputs[0]
# 将ResNet的输入切换为新定义的输入节点
ge.reroute_inputs(new_input.op, original_input_edge.op)

# 删除解码模块的所有节点（假设解码模块的节点都以'decode/'为前缀）
decode_module_ops = [op for op in tf.get_default_graph().get_operations() if op.name.startswith('decode/')]
ge.remove_nodes(decode_module_ops)

3. 导出修改后的冻结图 #

完成图结构修改后，将新图导出为冻结图文件：

# 获取新图的GraphDef结构
new_graph_def = tf.get_default_graph().as_graph_def()

# 移除训练相关的冗余节点
new_graph_def = tf.graph_util.remove_training_nodes(new_graph_def)

# 保存修改后的冻结图
with tf.gfile.GFile('modified_frozen_graph.pb', 'wb') as f:
    f.write(new_graph_def.SerializeToString())

额外说明 #

• 如果只有.ckpt文件，需要先加载.meta文件恢复图结构，再加载.ckpt的变量值，最后用tf.graph_util.convert_variables_to_constants将变量转为常量，导出为冻结图后再执行上述修改步骤。
• 确保新输入的维度、数据类型、预处理方式（比如归一化范围）和ResNet原输入完全一致，否则会出现维度不匹配或推理结果异常的问题。

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