我正在使用 Keras 和 TensorFlow 2.12.0 来实现具有自定义方法的子类模型,例如自定义
train_step
。为了保存模型,我使用 ModelCheckpoint
回调。
我的目标不仅是保存模型的架构和权重,还保存自定义方法和训练配置,使我能够从停止的确切状态恢复训练。从 Keras/TensorFlow 文档来看,
HDF5
格式似乎只能保存子类模型的权重。因此,我尝试使用 SavedModel
和 Keras
格式。但是加载模型后,训练配置和自定义方法不可用。
我可能遗漏了什么吗? 重现问题的工作代码片段:https://colab.research.google.com/drive/19tc4EY_wI66ZZBeLlhs0rqIyfQIKENhd?usp=sharing
经过一番调查,我意识到,由于子类化模型只是训练两个功能模型的包装,所以我需要:
我希望这可以帮助任何将来面临同样问题的人。