保存神经网络已训练权重用于其他数据集训练的TensorFlow实现问题

这个需求在两个版本下都有官方成熟实现，无需自定义复杂逻辑，具体操作如下：

TensorFlow 1.x 实现方式 #

步骤1：训练阶段保存最优权重 #

你可以直接用tf.train.Saver结合验证集指标判断触发保存，也可以用TF1.x兼容的Keras回调简化操作：

• 原生会话模式示例：

# 定义MLP架构、损失、优化器等逻辑
saver = tf.train.Saver()
best_val_acc = 0.0

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for epoch in range(epochs):
        # 训练步骤省略
        # 跑验证集获取准确率val_acc
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            saver.save(sess, "./ckpt/best_mlp.ckpt") # 保存最优权重

步骤2：加载权重作为新任务初始化 #

定义和训练时完全一致的MLP架构（变量名必须完全匹配），直接恢复权重即可：

# 定义和预训练阶段完全相同的MLP架构
saver = tf.train.Saver()

with tf.Session() as sess:
    saver.restore(sess, "./ckpt/best_mlp.ckpt") # 直接加载预训练权重，无需执行全局变量初始化
    # 接下来直接在新数据集上训练即可，需要冻层的话给对应变量设置trainable=False即可

TensorFlow 2.x 实现方式 #

2.x版本基于Keras高阶API实现更简单：

步骤1：训练阶段保存最优权重 #

直接用ModelCheckpoint回调自动保存验证集最优权重：

import tensorflow as tf

# 定义MLP模型
model = tf.keras.models.Sequential([
    # 这里填和你预训练一致的层结构，示例省略
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 配置最优权重保存回调
checkpoint_cb = tf.keras.callbacks.ModelCheckpoint(
    filepath="./best_mlp_weights.h5",
    monitor="val_accuracy",
    save_best_only=True,
    save_weights_only=True,
    verbose=1
)

# 训练时传入回调即可自动保存最优权重
model.fit(x_mnist, y_mnist, validation_split=0.1, epochs=20, callbacks=[checkpoint_cb])

步骤2：加载权重初始化新任务 #

定义同架构的模型后直接加载权重即可：

# 定义和预训练完全相同的MLP架构
new_model = tf.keras.models.Sequential([
    # 和预训练模型的层结构完全一致
])
new_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 加载预训练权重
new_model.load_weights("./best_mlp_weights.h5")

# 直接在新数据集上训练即可，需要微调部分层的话手动设置对应层的trainable属性为False即可

注意：两个版本下都需要保证加载权重的模型和保存权重的模型架构、变量/层名完全匹配，如果只需要加载部分层的权重，可以在加载时添加by_name=True参数，仅加载名称匹配的层的权重。

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