How to provide correct input for NMS layer

[afazel]

Describe the issue

I have an onnx format of complex-yolo pretrained model which has been created by this code:

torch.onnx.export(model,               # model being run
                input_imgs,                         # model input (or a tuple for multiple inputs)
                "./Complex_yolo_300.onnx",   # where to save the model (can be a file or file-like object)
                export_params=True,        # store the trained parameter weights inside the model file
                opset_version=11,          # the ONNX version to export the model to
                do_constant_folding=True,  # whether to execute constant folding for optimization
                input_names = ['input'],   # the model's input names
                output_names = ['output'], # the model's output names
               )

Now I want to load this model, add NMS layer to it and save it.I added the NMS layer by (the code is adopted from [here][1]:):

onnx_model = onnx.load("Complex_yolo_epoch_300.onnx")
    onnx.checker.check_model(onnx_model)
    onnx_model_with_nms_path = "Complex_yolo_epoch_300_with_nms.onnx"

    graph = onnx_model.graph

    # operation to transpose bbox before pass to NMS node
    transpose_scores_node = onnx.helper.make_node(
    'Transpose', 
    inputs=['output'], 
    outputs=['scores_transposed'],
    perm=(0, 2, 1))
    graph.node.append(transpose_scores_node)

    # make constant tensors for nms
    score_threshold = onnx.helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.25])
    iou_threshold = onnx.helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.45])
    max_output_boxes_per_class = onnx.helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [200])

    # create the NMS node
    inputs=["output",'scores_transposed' , 'max_output_boxes_per_class', 'iou_threshold', 'score_threshold',]
    # inputs=['onnx::Concat_458', 'onnx::Concat_459', 'max_output_boxes_per_class', 'iou_threshold', 'score_threshold',]
    outputs = ["selected_indices"]
    nms_node = onnx.helper.make_node(
        'NonMaxSuppression',
        inputs,
        outputs=["selected_indices"],
        center_point_box=1, 
    )

    # add NMS node 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=["num_results",3])
    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_model_with_nms_path)

By running this the model is created successfully, however when I want to test it on a sample image it gives me an error:

ort_session = onnxruntime.InferenceSession("Complex_yolo_epoch_300_with_nms.onnx")

    test_dataloader = create_test_dataloader(configs)
    with torch.no_grad():
        for batch_idx, (img_paths, imgs_bev) in enumerate(test_dataloader):
            input_imgs = imgs_bev.to(device).float()
            # compute ONNX Runtime output prediction
            ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_imgs)}
            ort_outs = ort_session.run(None, ort_inputs)

The error:

ort_outs = ort_session.run(None, ort_inputs)
  File "/home/azar/anaconda3/envs/pytorch_1_1/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 220, in run
    return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running NonMaxSuppression node. Name:'' Status Message: non_max_suppression.cc:92 PrepareCompute The most inner dimension in boxes must have 4 data.

I know that the shape of boxes for NMS layer should be [num_batches, spatial_dimension, 4] and in my case the output of complex_yolo model is [num_batches, spatial_dimension, 10], but I do not know how should I extract boxes and scores using the output of original onnx model, any suggestion?

To reproduce

ort_session = onnxruntime.InferenceSession("Complex_yolo_epoch_300_with_nms.onnx")

test_dataloader = create_test_dataloader(configs)
with torch.no_grad():
    for batch_idx, (img_paths, imgs_bev) in enumerate(test_dataloader):
        input_imgs = imgs_bev.to(device).float()
        # compute ONNX Runtime output prediction
        ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_imgs)}
        ort_outs = ort_session.run(None, ort_inputs)

Urgency

No response

Platform

Linux

OS Version

Ubuntu 22.4

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.14.1

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[hariharans29]

Hi @afazel - Unfortunately, it is hard to provide guidance for this without knowing what the output of your pre-trained model semantically means. If the output of your model indeed does have shape - [num_batches, spatial_dimension, 10] then it is hard to see how it can be bounding boxes given that bounding boxes are characterized by 4 values - Please see here.

Please consider asking if the model author can support the NMS layer in the model itself or atleast seek clarity as to how to extract bounding boxes from the model as only they can give you the right guidance for this as they know the semantics of the model.

Closing as this isn't really an ORT issue. Please feel free to open a discussion about this to see if other community members can share similar experiences if any.