Collectives™ on Stack Overflow

Find centralized, trusted content and collaborate around the technologies you use most.

Teams

Q&A for work

Connect and share knowledge within a single location that is structured and easy to search.

I have an ONNX model where I get the output as an 1xN array of probabilities. I want to add ArgMax to the end of the model so that I get the index instead.

I tried doing this using onnx.helper but was not able to find a good way to do it. I can create an ArgMax node using:

argmax = onnx.helper.make_node(
    'ArgMax',
    inputs=['inp'],
    outputs=['out'],
    axis=0,
    keepdims=0)
but how do I append this node to the end of the graph?
I usually use the onnx.compose.merge operator. If you have an existing model original_model that has some named output, in this example output_original_model you can create a graph from your node and use compose to combine them.
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")]
        
Thanks for contributing an answer to Stack Overflow!
Please be sure to answer the question. Provide details and share your research!
But avoid …
Asking for help, clarification, or responding to other answers.
Making statements based on opinion; back them up with references or personal experience.
To learn more, see our tips on writing great answers.