How to add a node to the end of a graph

I am using PyTorch to export a model, but the exporter does not export NMS. So one solution is for me to export before NMS and then add NMS node = onnx.helper.make_node('NonMaxSuppression') manually after loading the exported model.

[avasbr]

Got it. Thanks again for all the help @dashesy

As an aside, I dug a bit more into the onnx docs and found this to be quite useful: https://github.com/onnx/onnx/blob/master/docs/IR.md

It doesn’t cover specific examples of how to do things, but it clarified a number of concepts for me. In the end, I was able to attach a NonMaxSuppression op with score_threshold, iou_threshold, and max_output_boxes_per_class as constant values, and pass the onnx checker as well. For whoever it might help in the future, here’s the full code that worked for me, when going from an onnx model that currently outputs decoded bboxes and scores:

onnx version: 1.7.0 pytorch version: 1.4.0

onnx_fpath = [Path to .onnx model]
onnx_model = onnx.load(onnx_fpath)

# make constant tensors
score_threshold = onnx.helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.9])
iou_threshold = onnx.helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.3])
max_output_boxes_per_class = onnx.helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [200])

# create the NMS node
inputs=['bboxes', 'scores', 'max_output_boxes_per_class', 'iou_threshold', 'score_threshold',]
outputs = ["selected_indices"]
nms_node = onnx.helper.make_node(
    'NonMaxSuppression',
    inputs,
    ["selected_indices"],
)

# add to the list of graph nodes
graph.node.append(nms_node)

# append to the output (now the outputs would be scores, bboxes, selected_indices)
output_value_info = onnx.helper.make_tensor_value_info("selected_indices", TensorProto.INT64, shape=[])
graph.output.append(output_value_info)

# add to initializers - without this, onnx will not know where these came from, and complain that 
# they're neither outputs of other nodes, nor inputs. As initializers, however, they are treated 
# as constants needed for the NMS op
graph.initializer.append(score_threshold)
graph.initializer.append(iou_threshold)
graph.initializer.append(max_output_boxes_per_class)

# check that it works and re-save
onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, onnx_fpath)

[fumihwh]

First way: If you want to add a node to the end of a graph, use onnx.helper to make a node and append to model.graph.node is right way. Don’t forget to modify graph.output also. Second way: modify your code, add support to export nms to exporter.