如何加载TensorFlow模型续训，避免优化器状态重置？

核心问题在于：加载模型后重新执行model.compile()会重置优化器的训练状态（如动量、学习率衰减值等），导致续训时模型无法延续之前的训练节奏，验证损失出现跳升。以下是针对性解决方案：

方案1：保存并加载完整模型（推荐） #

训练时保存包含优化器状态的完整模型，加载后直接续训，无需重新编译：

• 训练结束后保存模型：

  # 保存完整模型（默认包含优化器状态、权重、模型结构）
  model.save('trained_model_full')
  

• 加载模型并直接续训：

  # 加载完整模型
  loaded_model = tf.keras.models.load_model('trained_model_full')
  
  # 直接启动续训，指定起始轮次
  loaded_model.fit(
      train_dataset,
      validation_data=val_dataset,
      epochs=50,  # 总训练轮次，比如之前训了20轮，这里设为50
      initial_epoch=20  # 从第20轮结束后开始续训
  )
  

方案2：单独保存权重与优化器状态（需自定义模型结构） #

如果只保存了权重文件，需要手动恢复优化器状态：

1. 训练时额外保存优化器状态：

   # 训练结束后保存模型权重
   model.save_weights('model_weights.h5')
   # 保存优化器权重（状态）
   import numpy as np
   np.save('optimizer_weights.npy', model.optimizer.get_weights())
   

2. 续训时恢复：

   # 重建与原训练完全一致的模型结构
   def build_model():
       # 这里写你原模型的构建代码
       inputs = tf.keras.Input(shape=(28,28))
       x = tf.keras.layers.Flatten()(inputs)
       x = tf.keras.layers.Dense(64, activation='relu')(x)
       outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
       return tf.keras.Model(inputs=inputs, outputs=outputs)
   
   model = build_model()
   # 编译模型，必须使用与原训练完全一致的优化器配置（学习率、动量等）
   optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
   model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
   
   # 加载模型权重
   model.load_weights('model_weights.h5')
   # 恢复优化器状态
   optimizer_weights = np.load('optimizer_weights.npy', allow_pickle=True)
   model.optimizer.set_weights(optimizer_weights)
   
   # 开始续训
   model.fit(...)
   

关键注意事项 #

• 优化器配置必须完全匹配：续训时使用的优化器类型、学习率、动量、衰减系数等参数，要和原训练时完全一致，否则状态恢复无效。
• 自定义组件需注册：如果模型使用了自定义损失函数、层或优化器，加载时要通过custom_objects参数注册，比如：

  loaded_model = tf.keras.models.load_model(
      'trained_model_full',
      custom_objects={'CustomLoss': CustomLoss, 'CustomLayer': CustomLayer}
  )
  

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