如何将 ArgMax 添加到 ONNX 模型的末尾？

【问题标题】：How to add ArgMax to the end of ONNX model?如何将 ArgMax 添加到 ONNX 模型的末尾？
【发布时间】：2022-11-07 10:58:12
【问题描述】：

我有an ONNX model，我将输出作为 1xN 概率数组。我想将 ArgMax 添加到模型的末尾，以便获得索引。

我尝试使用onnx.helper 执行此操作，但找不到一个好方法。我可以使用以下方法创建一个 ArgMax 节点：

argmax = onnx.helper.make_node(
    'ArgMax',
    inputs=['inp'],
    outputs=['out'],
    axis=0,
    keepdims=0)

但是如何将此节点附加到图表的末尾？

【问题讨论】：

检查这个towardsdatascience.com/merging-onnx-graphs-da088dbaf9b4

标签： onnx

【解决方案1】：

我通常使用onnx.compose.merge 运算符。如果您有一个现有模型 original_model 具有一些命名输出，在此示例中 output_original_model 您可以从您的节点创建一个图形并使用 compose 组合它们。

original_model = ...

argmax = onnx.helper.make_node(
    'ArgMax',
    inputs=['inp'],
    outputs=['out'],
    axis=0,
    keepdims=0)

graph = onnx.helper.make_graph(
    nodes=[
        argmax,
    ],
    name="argmaz",
    inputs=[
        onnx.helper.make_tensor_value_info(
            "inp", onnx.TensorProto.FLOAT, [None, None, input_size]
        )
    ],
    outputs=[
        onnx.helper.make_tensor_value_info(
            "out", onnx.TensorProto.FLOAT, [None, out_size]
        )
    ],
)

combined_model = onnx.compose.merge_models(
    onnx_original, graph, io_map=[("output_original_model", "inp")]
)

【讨论】：