AYE89

coding & learning

导航

Keras处理已保存模型中的自定义层(或其他自定义对象)

Posted on 2018-04-16 10:26  AYE89  阅读(1100)  评论(0编辑  收藏  举报

如果要加载的模型包含自定义层或其他自定义类或函数,则可以通过 custom_objects 参数将它们传递给加载机制:

from keras.models import load_model
# 假设你的模型包含一个 AttentionLayer 类的实例
model = load_model('my_model.h5', custom_objects={'AttentionLayer': AttentionLayer})

或者,你可以使用 自定义对象作用域

from keras.utils import CustomObjectScope

with CustomObjectScope({'AttentionLayer': AttentionLayer}):
    model = load_model('my_model.h5')

自定义对象的处理与 load_modelmodel_from_jsonmodel_from_yaml 的工作方式相同:

from keras.models import model_from_json
model = model_from_json(json_string, custom_objects={'AttentionLayer': AttentionLayer})