keras-BERT ValueError: Unknown layer: TokenEmbedding
问题复现:
训练的模型存储方式:
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)