如何将TensorFlow checkpoint导入Keras？训练前导入修改模型可行吗？

嘿，这个问题相当实用——不少人在跨框架迁移、微调模型时都会碰到这类需求。我分两部分给你讲得明明白白：

一、TensorFlow Checkpoint导入Keras的两种常见场景 #

场景1：Checkpoint是用tf.keras直接保存的 #

如果你的checkpoint是通过model.save_weights('my_model.ckpt')这种tf.keras原生方式保存的，那导入起来毫无难度：

# 第一步：先定义好和原模型结构完全一致的Keras模型
def build_original_model():
    inputs = tf.keras.Input(shape=(28,28))
    x = tf.keras.layers.Flatten()(inputs)
    x = tf.keras.layers.Dense(128, activation='relu')(x)
    outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
    return tf.keras.Model(inputs=inputs, outputs=outputs)

model = build_original_model()
# 第二步：直接加载权重
model.load_weights('path/to/my_model.ckpt')

这种情况下变量名完全匹配，不会有任何额外问题。

场景2：Checkpoint是原生TensorFlow（非tf.keras）保存的 #

如果是用tf.train.Checkpoint保存的原生TF模型权重，就得处理变量名映射的问题——原生TF和Keras的变量命名规则存在差异。

步骤如下：

1. 先查看checkpoint里的变量列表，确认每个变量的名称和形状：

# 列出checkpoint中的所有变量
from tensorflow.train import list_variables
vars_in_ckpt = list_variables('path/to/your/checkpoint-xxx')
for name, shape in vars_in_ckpt:
    print(f"变量名: {name}, 形状: {shape}")

2. 定义Keras模型后，通过tf.train.Checkpoint建立变量映射，再恢复权重：

keras_model = build_original_model()

# 用Checkpoint包装Keras模型，自动匹配结构相似的变量
checkpoint = tf.train.Checkpoint(model=keras_model)
# 恢复权重，expect_partial()用来忽略优化器等不需要的变量（如果存在）
status = checkpoint.restore('path/to/your/checkpoint-xxx').expect_partial()

# 可选：验证是否加载成功
status.assert_existing_objects_matched()

如果自动匹配失败，还可以手动指定变量映射：

checkpoint = tf.train.Checkpoint(
    dense_kernel=keras_model.layers[1].kernel,  # Keras的dense层权重
    dense_bias=keras_model.layers[1].bias,      # Keras的dense层偏置
    # 其他变量依次对应
)
checkpoint.restore('path/to/checkpoint-xxx').expect_partial()

二、训练前导入Checkpoint并修改模型：完全可行！ #

这其实是迁移学习、模型微调的常规操作，我自己做过很多次。核心思路是：先加载原模型权重，再基于原模型结构做修改，最后重新编译训练。

举个具体的例子：

# 1. 加载原模型权重
base_model = build_original_model()
base_model.load_weights('path/to/my_model.ckpt')

# 2. 修改模型：比如在原模型顶部新增分类层，适配新任务
x = base_model.output
x = tf.keras.layers.Dense(256, activation='relu', name='new_dense')(x)
new_output = tf.keras.layers.Dense(5, activation='softmax', name='new_output')(x)

# 构建修改后的模型
modified_model = tf.keras.Model(inputs=base_model.input, outputs=new_output)

# 3. 可选：冻结原模型的层，只训练新增的层（减少计算量）
for layer in base_model.layers:
    layer.trainable = False

# 4. 编译并开始训练
modified_model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)
modified_model.fit(train_dataset, epochs=10, validation_data=val_dataset)

注意事项 #

• 如果修改模型时删除了原模型的某些层，加载权重时可能会出现“未匹配变量”的警告，用expect_partial()可以忽略这些无关变量；
• 务必保证原模型和Keras模型的变量形状一致，否则会加载失败；
• 如果需要保留原训练的优化器状态（比如继续原训练进度），需要同时加载优化器的变量，但这种场景比较少见，大部分情况只需要加载模型权重即可。

内容的提问来源于stack exchange，提问作者D.Giunchi