复杂组合模型训练中优化器状态的保存与加载问题咨询

我太懂这种复杂模型训练时的糟心了——当模型是多个子模块拼出来的，还有权重共享、部分层冻结这些操作，常规的model.save()和load_model()确实会各种出问题。别慌，咱们拆成几个步骤来搞，把模型结构、权重、优化器状态分开保存加载，完美适配你的场景：

这种方式能避开复杂结构带来的序列化问题，精准控制每一部分的保存：

• 保存模型结构：如果你的模型有自定义层或者复杂组合，用序列化结构的方式比直接存整个模型靠谱。用to_json()（或者to_yaml()）把结构转成文本格式：

  # 保存结构到JSON文件
  model_structure_json = your_trainer_model.to_json()
  with open("trainer_model_structure.json", "w") as f:
      f.write(model_structure_json)
  

• 保存模型权重：用save_weights()单独存权重，不管是共享的权重还是冻结层的权重，都会被完整保存，而且不会带结构相关的冗余信息：

  # 保存权重到H5文件
  your_trainer_model.save_weights("trainer_model_weights.h5")
  

• 保存优化器状态：优化器的状态（比如动量值、学习率调度的进度）是恢复训练的关键，得单独存下来。我们可以把优化器的配置和权重一起用pickle序列化：

  import pickle
  
  # 提取优化器的配置和权重
  optimizer_config = your_trainer_model.optimizer.get_config()
  optimizer_weights = your_trainer_model.optimizer.get_weights()
  
  # 保存到pickle文件
  with open("optimizer_state.pkl", "wb") as f:
      pickle.dump((optimizer_config, optimizer_weights), f)
  

按照保存的逆顺序来，一步步恢复模型和训练状态：

• 重建模型结构：先从保存的JSON/YAML文件里加载结构，注意要确保你的自定义层、子模型类已经在当前环境中定义好了，不然会解析失败：

  from keras.models import model_from_json
  
  # 加载结构并重建模型
  with open("trainer_model_structure.json", "r") as f:
      loaded_structure = f.read()
  restored_model = model_from_json(loaded_structure)
  

• 加载模型权重：给重建好的模型加载之前保存的权重，只要结构和原来完全一致，共享权重会自动对应到正确的层上：

  # 加载权重
  restored_model.load_weights("trainer_model_weights.h5")
  

• 恢复优化器状态：先根据保存的配置重新实例化优化器，绑定到模型上后再加载权重，这样训练进度就完全接上了：

  from keras.optimizers import get
  
  # 加载优化器的配置和权重
  with open("optimizer_state.pkl", "rb") as f:
      opt_config, opt_weights = pickle.load(f)
  
  # 重建优化器
  restored_optimizer = get(opt_config)
  # 先编译模型，才能给优化器设置权重
  restored_model.compile(optimizer=restored_optimizer, loss=your_loss_func, metrics=your_metrics)
  # 加载优化器的状态权重
  restored_model.optimizer.set_weights(opt_weights)
  

• 权重共享的处理：只要重建的模型和原模型的结构完全一致（包括共享层的引用关系），加载权重时共享的权重会自动同步，不需要额外操作。
• 冻结层的保持：save_weights()只保存权重值，不会记录层的trainable属性，所以加载后原来冻结的层还是冻结状态。如果需要调整冻结状态，手动设置layer.trainable = True/False后重新编译模型就行。
• 子类化模型的情况：如果你的训练器模型是继承keras.Model的子类化模型，to_json()可能无法序列化。这时候直接手动创建一个模型实例（和原模型结构完全一样），然后直接调用load_weights()加载权重，优化器的保存加载逻辑和上面一致。

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