Keras子类化模型部署保存问题及解决方案咨询

针对Keras子类化模型无法直接保存的问题，这里有几个实用的解决方法：

1. 保存为TensorFlow SavedModel格式（推荐） #

子类化模型只要在训练或推理时处理过至少一次输入数据（你的代码中已经完成model.fit()训练，模型已获取输入形状信息），就可以直接用TensorFlow的SavedModel格式保存，这是官方推荐的部署友好格式。

# 训练完成后保存模型
model.save("transformer_savedmodel")

# 部署时加载模型
loaded_model = tf.keras.models.load_model("transformer_savedmodel")

2. 单独保存权重+重建模型结构 #

如果只需要保存模型参数，可以先保存权重文件，之后通过重新定义模型结构来加载权重：

# 保存权重
model.save_weights("transformer_weights.h5")

# 部署步骤：先重建完全一致的模型结构
loaded_model = Transformer(
    num_hid=200,
    num_head=2,
    num_feed_forward=400,
    target_maxlen=max_target_len,
    num_layers_enc=4,
    num_layers_dec=1,
    num_classes=34,
)
# 必须先让模型获取输入形状（用任意符合格式的输入调用一次）
loaded_model(batch[0])  # 假设batch[0]是你的输入数据
# 加载权重
loaded_model.load_weights("transformer_weights.h5")

3. 导出为ONNX格式（跨框架部署） #

如果需要在PyTorch、ONNX Runtime等非TensorFlow环境部署，可以先将SavedModel转换为ONNX格式：

import tf2onnx

# 先保存为SavedModel
model.save("transformer_savedmodel")
# 转换为ONNX文件
tf2onnx.convert.from_saved_model("transformer_savedmodel", output_path="transformer_model.onnx")

内容的提问来源于stack exchange，提问作者Omar Bouaziz