TensorFlow中如何向加载的Checkpoint添加新变量并复用子图权重

在TensorFlow 1.x中处理这种部分权重恢复的场景，核心思路是只恢复你需要的原有子图变量，而不是加载整个旧图。下面是几种实用的方法，根据你的场景选择：

方法1：保持变量名一致，批量恢复原有子图变量 #

这是最省心的方式，只要你在构建新计算图时，把原有子图的代码放在和保存时相同的variable_scope下，就能批量匹配变量名，快速恢复权重。

步骤： #

1. 构建新计算图时复用原有作用域
   把原来的子图代码包裹在和保存时完全一样的tf.variable_scope里，新增的网络放在新的作用域中：

   # 原有子图（和训练时结构完全一致，作用域名称相同）
   with tf.variable_scope("old_model"):
       x = tf.placeholder(tf.float32, shape=[None, 784])
       original_logits = your_original_model(x)
   
   # 新增的神经网络部分
   with tf.variable_scope("new_addon"):
       final_output = your_new_addon_model(original_logits)
   

2. 创建仅恢复原有变量的Saver
   通过作用域筛选出原有子图的所有变量，用这些变量创建Saver，再执行恢复操作：

   def load_original_subgraph_weights(sess, param_folder, saved_ckpt):
       print("Loading weights for original subgraph...")
       ckpt_path = os.path.join(param_folder, 'model', saved_ckpt)
       
       # 筛选出原有作用域下的所有全局变量
       original_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="old_model")
       # 只针对这些变量创建Saver
       saver = tf.train.Saver(var_list=original_vars)
       
       saver.restore(sess, ckpt_path)
       print("Original subgraph weights restored successfully!")
   

方法2：变量名映射（当原有变量名有改动时） #

如果因为重构，新图中原有子图的变量名/作用域变了，可以手动构建旧变量名→新变量对象的映射字典，让TensorFlow正确匹配权重。

示例代码： #

def load_with_name_mapping(sess, param_folder, saved_ckpt):
    print("Loading weights with variable name mapping...")
    ckpt_path = os.path.join(param_folder, 'model', saved_ckpt)
    
    # 获取新图中原有子图的变量（假设新作用域是"legacy_model"）
    new_original_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="legacy_model")
    
    # 构建映射：Checkpoint中的旧变量名 → 新图中的变量对象
    name_mapping = {
        "old_model/dense1/weights": new_original_vars[0],
        "old_model/dense1/biases": new_original_vars[1],
        "old_model/dense2/weights": new_original_vars[2],
        # 按Checkpoint中的变量名依次对应
    }
    
    saver = tf.train.Saver(var_list=name_mapping)
    saver.restore(sess, ckpt_path)
    print("Mapped weights restored successfully!")

小技巧：查看Checkpoint中的变量名 #

如果你不确定Checkpoint里的变量名，可以用以下代码打印所有变量信息：

from tensorflow.python.tools.inspect_checkpoint import print_tensors_in_checkpoint_file

ckpt_path = os.path.join(param_folder, 'model', saved_ckpt)
print_tensors_in_checkpoint_file(
    file_name=ckpt_path,
    tensor_name='',  # 空字符串表示打印所有变量
    all_tensors=False,
    all_tensor_names=True
)

方法3：用init_from_checkpoint灵活初始化 #

TensorFlow提供了tf.train.init_from_checkpoint函数，可以直接指定从Checkpoint初始化哪些变量，不需要创建Saver，适合在初始化阶段批量处理。

示例代码： #

def setup_checkpoint_initialization(param_folder, saved_ckpt):
    ckpt_path = os.path.join(param_folder, 'model', saved_ckpt)
    
    # 方式1：整个作用域映射（旧作用域→新作用域）
    tf.train.init_from_checkpoint(ckpt_path, {
        "old_model/": "old_model/"  # 旧图中old_model下的变量，对应新图中old_model下的变量
    })
    
    # 方式2：单个变量映射（适合部分变量恢复）
    # tf.train.init_from_checkpoint(ckpt_path, {
    #     "old_model/dense1/weights": "legacy_model/dense1/weights",
    #     "old_model/dense1/biases": "legacy_model/dense1/biases"
    # })

# 先构建计算图
setup_checkpoint_initialization(param_folder, saved_ckpt)

# 初始化所有变量：原有子图变量从Checkpoint加载，新增变量用默认初始化
init_op = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init_op)
    # 模型可以正常使用了

关键注意事项 #

• 原有子图的变量形状和数据类型必须和Checkpoint中保存的完全一致，否则会恢复失败。
• 不要用tf.train.import_meta_graph加载旧图，因为你需要的是把原有子图整合到新图中，而不是替换整个图。
• 如果你的新图在原有子图基础上做了小修改（比如新增了层），只要原有变量的结构没改，依然可以用上述方法恢复原有权重。

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