TensorFlow冻结图tf.const转tf.variable重训练遇类型兼容及初始化问题

报错原因说明 #

• 第一个ValueError：TensorFlow 2.x中tf.Variable默认返回资源类型张量，直接用变量节点输出替换Const节点的浮点型输出，数据类型不匹配。你后续调整为替换ReadVariableOp节点输出的逻辑是正确的。
• 第二个FailedPreconditionError：通过Graph Editor修改图结构后，新生成的变量没有被纳入当前TensorFlow的变量追踪体系，也没有被显式初始化，导致推理时找不到对应的资源句柄。

修复后完整代码 #

import tensorflow as tf
import tensorflow.contrib.graph_editor as ge

def convert_frozen_graph_to_trainable(graph_def):
    # 导入原始冻结图定义
    train_graph = tf.Graph()
    with train_graph.as_default():
        tf.graph_util.import_graph_def(graph_def, name='')
    
    const_to_var_map = []
    converted_vars = []
    
    with train_graph.as_default():
        # 复用同一个Session读取所有Const节点的值，减少性能损耗
        with tf.compat.v1.Session() as sess:
            const_ops = [op for op in train_graph.get_operations() if op.type == "Const"]
            const_tensor_list = [train_graph.get_tensor_by_name(f"{op.name}:0") for op in const_ops]
            const_value_list = sess.run(const_tensor_list)
            
            for const_op, np_val in zip(const_ops, const_value_list):
                var_name = f"{const_op.name}_turned_var"
                # 创建可训练变量，直接传入numpy值作为初始值
                new_var = tf.Variable(
                    name=var_name,
                    dtype=np_val.dtype,
                    shape=np_val.shape,
                    initial_value=np_val,
                    trainable=True
                )
                converted_vars.append(new_var)
                const_to_var_map.append((const_op.name, var_name))
    
    # 替换图中的Const节点为Variable的Read输出
    ge_graph = ge.Graph(train_graph.as_graph_def())
    for const_name, var_name in const_to_var_map:
        src_const_op = ge_graph._node_name_to_node[const_name]
        target_read_op = ge_graph._node_name_to_node[f"{var_name}/Read/ReadVariableOp"]
        ge.swap_outputs(ge.sgv(src_const_op), ge.sgv(target_read_op))
    
    # 转换回TensorFlow原生图并初始化变量
    final_train_graph = ge_graph.to_tf_graph()
    with final_train_graph.as_default():
        # 显式初始化所有转换得到的变量，写入当前资源容器
        var_init_op = tf.compat.v1.variables_initializer(converted_vars)
        with tf.compat.v1.Session() as sess:
            sess.run(var_init_op)
            # 若需要保存为SavedModel可在此处执行保存逻辑
            # tf.saved_model.save(sess, export_dir, signatures={...})
    
    return final_train_graph, converted_vars

额外注意事项 #

• TensorFlow 2.1.0版本的graph_editor属于contrib模块，本身兼容性有限，若条件允许可升级到更高版本，使用tensorflow-model-optimization工具集完成类似转换需求。
• 后续训练时需要将converted_vars传入优化器的var_list参数，否则变量不会参与梯度更新。
• 推理前可先执行一次变量初始化操作，避免资源容器找不到变量的报错。

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