在SageMaker Studio中将Keras模型转ONNX遇错的解决方法

你在SageMaker Studio中执行以下Keras模型转ONNX的代码时：

tf2onnx.convert.from_keras(model_f, 
                           output_path="2024-08-22-autocat-default.onnx", 
                           opset=14, 
                           extra_opset=[helper.make_opsetid('ai.onnx.contrib', 1)]);

出现错误：

AttributeError: 'Functional' object has no attribute '_get_save_spec'

但相同代码和依赖版本在SageMaker Jupyter Lab中运行正常，可尝试以下解决方法：

方法1：显式指定输入签名 #

SageMaker Studio环境中，tf2onnx可能需要明确的输入张量规范才能正确转换Functional模型。根据你的模型输入形状修改以下代码：

import tensorflow as tf
import tf2onnx
from tf2onnx import helper

# 替换为你的模型实际输入形状、数据类型和名称
input_signature = [tf.TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name="input")]

tf2onnx.convert.from_keras(model_f, 
                           output_path="2024-08-22-autocat-default.onnx", 
                           opset=14, 
                           input_signature=input_signature,
                           extra_opset=[helper.make_opsetid('ai.onnx.contrib', 1)])

方法2：先保存为SavedModel再转换 #

绕过直接转换Keras模型的环境差异，先将模型保存为TensorFlow标准的SavedModel格式，再转ONNX：

# 临时保存模型
model_f.save("./temp_saved_model")

# 从SavedModel转换为ONNX
tf2onnx.convert.from_saved_model("./temp_saved_model",
                                 output_path="2024-08-22-autocat-default.onnx",
                                 opset=14,
                                 extra_opset=[helper.make_opsetid('ai.onnx.contrib', 1)])

方法3：重新确认并修复依赖匹配 #

虽然你说版本一致，但可在两个环境中执行以下命令再次核对TensorFlow和tf2onnx的版本：

pip show tensorflow tf2onnx

若版本确实匹配，尝试重新安装tf2onnx以修复可能的环境缓存问题：

pip uninstall -y tf2onnx && pip install tf2onnx

方法4：移除非必需的扩展算子集 #

如果你的模型不需要使用ai.onnx.contrib中的特殊算子，可去掉extra_opset参数，简化转换命令：

tf2onnx.convert.from_keras(model_f, 
                           output_path="2024-08-22-autocat-default.onnx", 
                           opset=14)

内容的提问来源于stack exchange，提问作者Feng Chen