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': , 'MultiHeadAttention': , 'FeedForward': , 'TrigPosEmbedding': , 'EmbeddingRet': , 'EmbeddingSim': , 'PositionEmbedding': , 'TokenEmbedding': , 'EmbeddingSimilarity': , 'Masked': , 'Extract': , 'gelu': , 'gelu_tensorflow': , 'gelu_fallback': , 'AdamWarmup': }
当然,如果有些自定义的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)
 

你可能感兴趣的:(深度学习,keras,bert)