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

Question

我有一个通过模型子类编写的tensorflow模型，我想将它导出为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")],

如这笔记本所示