问题描述
我正在尝试创建一个超级简单的 Tensorflow 网络(一个数据处理 Lambda 层),然后将模型转换为 ONNX 并在从 onnxruntime 调用 ONNX 模型时验证结果匹配。我正在使用 Tensorflow v2.5.0。 & onnxruntime v1.8.1。
example_input2 = tf.convert_to_tensor([0,1000,2000,3000,4000],dtype=tf.float32)
模型定义:
inp = keras.layers.Input(name="input",type_spec=tf.TensorSpec(shape=[None],dtype=tf.float32))
output = keras.layers.Lambda(lambda x: tf.roll(x,shift=-1,axis=0),name="output") (inp)
model = keras.Model(inp,output,name="pipeline")
然后我可以将我的 example_input2 输入到网络中:
model.predict(example_input2)
提供所需的输出(简单的 tf.roll 操作):
array([1000.,2000.,3000.,4000.,0.],dtype=float32)
太好了!现在我可以保存我的 tensorflow 模型了,
model.save("OnnxInvestigateData/pipeline2",overwrite=True,include_optimizer=False,save_format='tf')
然后在 shell 中,我可以使用 tf2onnx 将其转换为 ONNX 格式:
python -m tf2onnx.convert --opset 14 --saved-model pipeline2 --output pipeline2.onnx
然后,回到 python,我可以加载 onnx 模型并尝试输入相同的输入:
sess = rt.InferenceSession("OnnxInvestigateData/pipeline2.onnx",log_verbosity_level=2)
xinput = example_input2.numpy()
sess.run(['output'],{"args_0":xinput})
提供与输入匹配的输出,而不是所需的输出(应该由 -1 tf.roll'd):
[array([ 0.,1000.,4000.],dtype=float32)]
当我在原始 keras 模型上从 python 内部调用 model.predict
时,我完全不知道为什么这里的输出不匹配。有什么想法吗?
解决方法
您似乎在 tf2onnx 中发现了一个错误。这是一个修复: https://github.com/onnx/tensorflow-onnx/pull/1616
如果您只是在进行测试并且不想等待修复程序合并,请尝试使用正移位值。
output = keras.layers.Lambda(lambda x: tf.roll(x,shift=2,axis=0),name="output") (inp)
另外,如果你想在你的脚本中直接从 keras 转换,你可以这样做:
onnx_model,_ = tf2onnx.convert.from_keras(model,opset=14)
sess = rt.InferenceSession(onnx_model.SerializeToString())