Modify the onnx model output example

foreword

insert image description here
The picture shows the onnx model opened in the netron (github link) software. You can see that the final output of the model on the right is the classification value predict_0 instead of the probability value. So how to obtain the probability value of the intermediate process, or how to cut off a part of the right picture and turn it into the left picture ?

the code

read model

import onnx
onnx_model = onnx.load("xxx.onnx")
graph = onnx_model.graph

First read your model in the form of a graph. The graph generally includes four parts: node (node), initializer (initialization), input (input), and output (output). It is very long to print out all of them. Here we mainly involve deleting nodes and modifying output .

Check output
insert image description here
to print output, and find that the main information is name and elem_type.
Use netron to select the upper layer node, namely identity, to view detailed information.

You can see that the output name of the upper layer node is the output name of output, and the two need to be consistent.
At the same time, elem_type=7 means that the output type is int64 (refer to my previous article), which also needs to be consistent with the output type of the upper layer node.

delete node node

nodes = graph.node
for i in range(len(nodes)):
    print(i,nodes[i])

Since it is modifying the output, we only care about the last few nodes.
insert image description here
Needless to say, the argmax function is a classic function converted to 0,1 output, so the probability value must be taken before argmax, and mul is multiplication. We need to use netron to open the model to view the mul node as follows: multiplying the 0,1 matrix should be converted to 2 columns for subsequent processing, then we only need to take the value of A, which is "add_result_0" output by the upper layer node
insert image description here
.
That is, to delete nodes with sequence numbers 261,262,263,264:

graph.node.remove(nodes[264])
graph.node.remove(nodes[263])
graph.node.remove(nodes[262])
graph.node.remove(nodes[261])

Delete in reverse order, and print the nodes again:
insert image description here
complete deletion.

modify output

graph.output[0].name = 'add_result_0'
graph.output[0].type.tensor_type.elem_type = 1

Since the output value of add_result_0 is float32, it is necessary to modify the elem_type in the output to the corresponding data type, 1 is float32, and see my previous article for more types.

Result comparison

onnx.save(onnx_model, 'modify_xxx.onnx')

Save the model first, and then call the onnx runtime to call the model:

import onnxruntime as rt
sess = rt.InferenceSession("modify_xxx.onnx")
onnx_pred = sess.run() #具体里面填什么根据你的模型填，此处为伪代码

insert image description here
Compared with the unmodified model:

it is found that the modified model has a probability value greater than 0 and is classified as 1, and that is less than 0 is classified as 0, indicating that our modification has been successful.