保存含tensorflow_hub层的Keras模型报NotImplementedError如何处理

错误原因 #

HDF5（.h5）格式保存Keras模型时，会序列化每一层的配置信息。hub.KerasLayer仅支持传入字符串类型的handle参数（可以是TF Hub模型标识、本地SavedModel文件夹路径）时生成可序列化的配置，你当前代码提前调用hub.load()得到模型对象再传入hub.KerasLayer，这类加载完成的模型对象无法被序列化，因此抛出NotImplementedError。 #

修复方案 #

方案1：修改KerasLayer的传参，适配HDF5格式保存 #

根据你是否需要微调TF Hub模型，选择对应写法：

• 无需微调TF Hub层
  直接将模型标识字符串传入hub.KerasLayer，无需提前加载：

# 删掉原来的loaded_obj = hub.load(module_url) 这行
emb = hub.KerasLayer(module_url, trainable=False)

• 需要微调TF Hub层
  微调完成后先将TF Hub模型对象保存为本地SavedModel，再将本地路径字符串传入hub.KerasLayer：

# 微调完成后先保存hub模型
tf.saved_model.save(loaded_obj, "./fine_tuned")
# 重新构建模型时传入本地路径字符串
emb = hub.KerasLayer("./fine_tuned", trainable=False)

修改后即可正常用model.save("model.h5", include_optimizer=False)保存模型，加载时需要指定自定义层：

model = tf.keras.models.load_model("model.h5", custom_objects={"KerasLayer": hub.KerasLayer})

方案2：改用TensorFlow原生SavedModel格式保存 #

直接将模型保存为SavedModel文件夹格式，不需要修改原有代码，该格式会直接保存完整计算图，无需序列化层配置，适配所有场景：

# 直接保存为文件夹，不需要.h5后缀
model.save("my_saved_model", include_optimizer=False)
# 加载时直接读取文件夹即可
model = tf.keras.models.load_model("my_saved_model")

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