问题复现:
训练的模型存储方式:
model.save_weights(model_path)
模型加载方式
from keras.models import load_model
model = load_model(model_path)
异常:
ValueError: Unknown layer: TokenEmbedding
出现该错误是因为要保存的model中包含了自定义的层(Custom Layer),导致加载模型的时候无法解析该Layer
解决该问题的方法是在load_model函数中添加custom_objects参数,该参数接受一个字典,键值为自定义的层,当然也可以偷懒,直接使用custom_objects=get_custom_objects()
解决方案:
from keras.models import load_model
from keras_bert import get_custom_objects
model = load_model(model_path, custom_objects=get_custom_objects())
keras_bert中提供的get_custom_objects() 中包含如下元素:
{'LayerNormalization':
当然,如果有些自定义的object不在上面这里面的话,那么需要将自定义的object放到里面,如下所示:
from keras.models import load_model
custom_objects = get_custom_objects()
my_objects = {'RAdam': RAdam}
custom_objects.update(my_objects)
model = load_model(model_path, custom_objects=custom_objects)