Keras NLP中TransformerDecoder模型保存失败：非JSON可序列化参数类型错误

问题代码 #

from keras_nlp.layers import TransformerDecoder
from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Embedding, GlobalAveragePooling1D, Input

decoder_input = Input(shape=(128,), name="decoder_input")
embedding = Embedding(input_dim=10000, output_dim=32, name="embedding")(decoder_input)
decoder = TransformerDecoder(intermediate_dim=64, num_heads=2, name="decoder")(embedding)
pooled = GlobalAveragePooling1D(name="pooling")(decoder)
output = Dense(token_size, activation="softmax", name="output_dense")(pooled)
model = Model(inputs=decoder_input, outputs=output)

model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])

model.save("model.keras")

错误信息 #

TypeError: Layer tf.__operators__.add was passed non-JSON-serializable arguments. 
Arguments had types: {'y': <class 'keras.src.backend.tensorflow.core.Variable'>, 'name': <class 'NoneType'>}. 
They cannot be serialized out when saving the model.

使用版本 #

print(tensorflow.__version__)
print(keras_nlp.__version__)
# 2.17.0
# 0.14.1

解决方法 #

• 升级Keras NLP版本：该序列化bug在Keras NLP 0.15.0及以上版本已修复，执行命令升级：

  pip install --upgrade keras-nlp>=0.15.0
  

• 临时替代方案（无法升级时）：改用TensorFlow SavedModel格式保存模型，替换原保存代码为：

  model.save("saved_model", save_format="tf")
  

  加载模型时执行：

  from tensorflow.keras.models import load_model
  model = load_model("saved_model", custom_objects={"TransformerDecoder": TransformerDecoder})
  

内容的提问来源于stack exchange，提问作者E.K.