Make tiny yolo v4 onnx output compatible with nms layer batchedNMSPlugin (C++/TensoRT)

Description

Hi all. From my C++/TensorRT code, I import a tiny yolo v4 ONNX, and I need to adapt the network output tensor to the required nms layer input tensor shapes (boxes input + scores input (according to TensorRT/plugin/batchedNMSPlugin at master · NVIDIA/TensorRT · GitHub )

What i get as network output:
B, nb_boxes_per_cell * box_params_size, Cx, Cy (1, 3*(4+1+80), 13, 13 = 1, 255, 13, 13)

The tensor shapes required by the batchedNMSPlugin layer:
→ 1, 255, 80, 4 (boxes input) (*)
→ 1, 255, 80 (scores input) (**)

What Ive done

1, 255, 13, 13 → 1, 255, 1313 (using setReshapeDimensions of a 1st IShuffleLayer)
1, 255, 13
13 → 1, 85, 13133 (using setReshapeDimensionsa of 2nd IShuffleLayer)
1, 85, 13133 → 1, 13133, 85 (using setSecondTranspose of the 2nd IShuffleLayer)

But now, how to reach (*) and (**)? Using the IGatherLayer? And in that case, how to define the dim indexes ITensor with constant values ?

Is there a alternative way to add nms to an imported tiny yolov4 onnx model (C++) ?

Please… help !

++

Hi,
Please refer to below links related custom plugin implementation and sample:

While IPluginV2 and IPluginV2Ext interfaces are still supported for backward compatibility with TensorRT 5.1 and 6.0.x respectively, however, we recommend that you write new plugins or refactor existing ones to target the IPluginV2DynamicExt or IPluginV2IOExt interfaces instead.

Thanks!

CORRECTED POST ==>

Hi all. From my C++/TensorRT code, I import a tiny yolo v4 ONNX, and I need to adapt the network output tensor to the required nms layer input tensor shapes (boxes input + scores input (according to TensorRT/plugin/batchedNMSPlugin at master · NVIDIA/TensorRT · GitHub )

What i get as network output:
B, nb_boxes_per_cell * box_params_size, Cx, Cy (1, 3*(4+1+80), 13, 13 = 1, 255, 13, 13)

The tensor shapes required by the batchedNMSPlugin layer:
→ 1, 13x13x3, 80, 4 (boxes input) (*)
→ 1, 13x13x3, 80 (scores input) (**)

What Ive done

1, 255, 13, 13 → 1, 255, 13x13 (using setReshapeDimensions of my 1st IShuffleLayer)
1, 255, 13x13 → 1, 255/3, 13x13x3 (using setReshapeDimensions of my 2nd IShuffleLayer)
1, 255/3, 13x13x3 → 1, 13x13x3, 255/3 (using setSecondTranspose of the 2nd IShuffleLayer)

But now, how to reach (*) and (**)? Using the IGatherLayer? And in that case, how to define the dim indexes ITensor with constant values ?

Is there a alternative way to add nms to an imported tiny yolov4 onnx model (C++) ?

Please… help !

++