将tensorflow模型导出到ONNX并指定变量名

Question

我有一个通过模型子类化编写的张量流模型，我想将其导出为 ONNX 格式。这很简单，附上脚本。但是，ONNX 模型的输入变量名称是 args_0。我如何重命名它？

import tensorflow as tf
import tf2onnx
from tensorflow.python.keras import Model
from tensorflow.python.keras.layers import Dense


class MyModel(Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dense = Dense(16)

    def call(self, inputs, **kwargs):
        return self.dense(inputs)

    def to_onnx(self, output_path, opset=14):
        model_proto, _ = tf2onnx.convert.from_keras(
            self,
            input_signature=[tf.TensorSpec((1, 128))],
            opset=opset,
            output_path=output_path,
        )
        return


if __name__ == "__main__":
    output_path = "./test.onnx"
    A = MyModel()
    A.to_onnx(output_path)

Answer 1

你可以在 input_signature 中提供输入名称 name="input_name" 所以，它应该看起来像

 input_signature=[tf.TensorSpec((1, 128), name="input_name")],

如图this笔记本

将tensorflow模型导出到ONNX并指定变量名

Export tensorflow model to ONNX and specify variable names

keras

tensorflow

onnx

tf.keras

tf2onnx