How to use NMS with Pytorch model (that was converted to ONNX -> TensorRT)

[ivanpanshin]

All right, so, I have a PyTorch detector SSD with MobileNet. Since I failed to convert model with NMS in it (to be more precise, I converted it, but TRT engine is built in a wrong way with that .onnx file), I decided to leave NMS part to TRT.

In general, there are several ways to add NMS in TRT:

1. Use graphsurgeon with TensorFlow model and add NMS as graphsurgeon.create_plugin_node
2. Use CPP code for plugin (https://github.com/NVIDIA/TensorRT/tree/master/plugin/batchedNMSPlugin)
3. Use DeepStream that has NMS plugin

But, I have a PyTorch model that I converted to onnx and then to TRT without any CPP code (Python only). My question is very simple: how can I combine my current pipeline with the CPP plugin for NMS?

[pranavm-nvidia]

You can use ONNX-GraphSurgeon to modify the ONNX model to include a plugin node. You can look at the onnx_packnet example for details.

[ivanpanshin]

Wow, ONNX-GS, nice! And just to be 100% sure: it's possible to use ONNX-GS in the same way as the regular GS to use this CPP NMS plugin (https://github.com/NVIDIA/TensorRT/tree/master/plugin/batchedNMSPlugin), correct?

[pranavm-nvidia]

Yes, correct. You can replace the old NMS subgraph/node with a new node whose op is set to the plugin name. There are two ways to do this, which I've outlined below.

Old-GS Style

In the simplest case, you can create the node and insert it into the graph like old GS, e.g.:

# Find tensors
tmap = graph.tensors()
boxes, scores, nms_out = tmap["boxes"], tmap["scores"], tmap["nms_out"]

# Disconnect old subgraph
boxes.outputs.clear()
scores.outputs.clear()
nms_out.inputs.clear()

attrs = {
    "share_location": False,
    "num_classes": 8,
    # etc.
}
node = gs.Node(op="BatchedNMS_TRT", attrs=attrs, 
               inputs=[boxes, scores], outputs=[nms_out])
graph.nodes.append(node)

Layer + Register API

If you're reusing this across multiple models, or using the plugin multiple times in the model, then the layer/register API might simplify things.

First you'd register a function to insert the plugin:

@gs.Graph.register()
def trt_batched_nms(self, boxes_input, scores_input, nms_output, 
                    share_location, num_classes): # and other attrs
    boxes_input.outputs.clear()
    scores_input.outputs.clear()
    nms_output.inputs.clear()

    attrs = {
        "share_location": share_location,
        "num_classes": num_classes,
        # etc.
    }
    return self.layer(op="BatchedNMS_TRT", attrs=attrs, 
                      inputs=[boxes_input, scores_input], 
                      outputs=[nms_output])

And then in your models, you can use it without all the boilerplate:

tmap = graph.tensors()

# Can also get attributes from the original graph instead of hard-coding
graph.trt_batched_nms(tmap["boxes"], tmap["scores"], 
                      tmap["nms_out"], share_location=True, 
                      num_classes=81)

graph.trt_batched_nms(tmap["boxes2"], tmap["scores2"], 
                      tmap["nms_out2"], share_location=False, 
                      num_classes=4)

[qraleq]

Hi, @pranavm-nvidia, I tried using the ONNX-GS method as you proposed and I've successfully replaced the NMS operator with BatchedNMS_TRT using the following code:

import onnx_graphsurgeon as gs
import onnx
import numpy as np

input_model_path = "model.onnx"
output_model_path = "model_gs.onnx"

@gs.Graph.register()
def trt_batched_nms(self, boxes_input, scores_input, nms_output,
                    share_location, num_classes):

    boxes_input.outputs.clear()
    scores_input.outputs.clear()
    nms_output.inputs.clear()

    attrs = {
        "shareLocation": share_location,
        "numClasses": num_classes,
        "backgroundLabelId": 0,
        "topK": 116740,
        "keepTopK": 100,
        "scoreThreshold": 0.3,
        "iouThreshold": 0.6,
        "isNormalized": True,
        "clipBoxes": True
    }
    return self.layer(op="BatchedNMS_TRT", attrs=attrs,
                      inputs=[boxes_input, scores_input],
                      outputs=[nms_output])


graph = gs.import_onnx(onnx.load(input_model_path))
graph.inputs[0].shape=[1,1280,720,3]
print(graph.inputs[0].shape)

for inp in graph.inputs:
    inp.dtype = np.int

input = graph.inputs[0]

tmap = graph.tensors()

graph.trt_batched_nms(tmap["Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_2/NonMaxSuppressionV5__1712:0"],
                      tmap["Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores/NonMaxSuppressionV5__1761:0"],
                      tmap["NonMaxSuppression__1763:0"],
                      share_location=False,
                      num_classes=4)

graph.trt_batched_nms(tmap["Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_2/NonMaxSuppressionV5__1712:0"],
                      tmap["Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_1/NonMaxSuppressionV5__1737:0"],
                      tmap["NonMaxSuppression__1739:0"],
                      share_location=False,
                      num_classes=4)

graph.trt_batched_nms(tmap["Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_2/NonMaxSuppressionV5__1712:0"],
                      tmap["Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_2/NonMaxSuppressionV5__1713:0"],
                      tmap["NonMaxSuppression__1715:0"],
                      share_location=False,
                      num_classes=4)

graph.trt_batched_nms(tmap["Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_2/NonMaxSuppressionV5__1712:0"],
                      tmap["Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_3/NonMaxSuppressionV5__1689:0"],
                      tmap["NonMaxSuppression__1691:0"],
                      share_location=False,
                      num_classes=4)


# Remove unused nodes, and topologically sort the graph.
# graph.cleanup()
# graph.toposort()
# graph.fold_constants().cleanup()

# Export the ONNX graph from graphsurgeon
onnx.checker.check_model(gs.export_onnx(graph))
onnx.save_model(gs.export_onnx(graph), output_model_path)

print("Saving the ONNX model to {}".format(output_model_path))

The problem I'm currently facing is that onnx.checker.check_model(gs.export_onnx(graph)) function returns following error:

onnx.onnx_cpp2py_export.checker.ValidationError: 
Node (NonMaxSuppression__1763) has output size 0 not in range [min=1, max=1].

==> Context: Bad node spec: 
input: "const_fold_opt__2119" 
input: "Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_3/iou_threshold:0" 
input:"Postprocessor/BatchMultiClassNonMaxSuppression/MultiClassNonMaxSuppression/non_max_suppression_with_scores_2/score_threshold:0" 
name: "NonMaxSuppression__1763" op_type: "NonMaxSuppression"

---

I think that the problem is in builtin_op_importers.cpp, where I added this code similar to #523:

{
DEFINE_BUILTIN_OP_IMPORTER(NonMaxSuppression)
    // NonMaxSuppression is not supported opset below 10.
    ASSERT(ctx->getOpsetVersion() >= 10, ErrorCode::kUNSUPPORTED_NODE);

    nvinfer1::ITensor* boxes_tensor = &convertToTensor(inputs.at(0), ctx);
    nvinfer1::ITensor* scores_tensor = &convertToTensor(inputs.at(1), ctx);
    const int numInputs = inputs.size();
    LOG_ERROR("no of inputs are "<<numInputs);
    LOG_ERROR("node outsize and op type are "<<node.output().size()<< " type " << node.op_type());

    const auto scores_dims = scores_tensor->getDimensions();
    const auto boxes_dims = boxes_tensor->getDimensions();
    
    LOG_ERROR("boxes dims "<< boxes_dims.nbDims << " dim3 has size "<<boxes_dims.d[2]);
    const std::string pluginName = "BatchedNMS_TRT";
    const std::string pluginVersion = "1";
    std::vector<nvinfer1::PluginField> f;

    /*
    bool share_location = true;
    const bool is_normalized = true;
    const bool clip_boxes = true;
    int backgroundLabelId = 0;
    
    // Initialize.
    f.emplace_back("shareLocation", &share_location, nvinfer1::PluginFieldType::kINT8, 1);
    f.emplace_back("isNormalized", &is_normalized, nvinfer1::PluginFieldType::kINT8, 1);
    f.emplace_back("clipBoxes", &clip_boxes, nvinfer1::PluginFieldType::kINT8, 1);
    f.emplace_back("backgroundLabelId", &backgroundLabelId, nvinfer1::PluginFieldType::kINT32, 1);
    */
    // Create plugin from registry
    //nvinfer1::IPluginV2* plugin = importPluginFromRegistry(ctx, pluginName, pluginVersion, node.name$
    nvinfer1::IPluginV2* plugin = createPlugin(node.name(), importPluginCreator(pluginName, pluginVers$

    ASSERT(plugin != nullptr && "NonMaxSuppression plugin was not found in the plugin registry!",
        ErrorCode::kUNSUPPORTED_NODE);

    std::vector<nvinfer1::ITensor*> nms_inputs ={boxes_tensor, scores_tensor};
    RETURN_FIRST_OUTPUT(ctx->network()->addPluginV2(nms_inputs.data(), nms_inputs.size(), *plugin));
}

Do you have an example of properly registering BatchedNMS_TRT as a plugin?
Also, do you have end to end example of replacing BatchedNMS with a TRT version?

[pranavm-nvidia]

@qraleq You should not need to modify the parser at all in TRT 7.1 and later (for any plugin, not just NMS). Modifying the node in the ONNX model should be enough.

Since BatchedNMS_TRT is not technically a valid ONNX op, it's expected that ONNX checker will fail.
From the error message though, it seems like you might have missed one?

name: "NonMaxSuppression__1763" op_type: "NonMaxSuppression"

Do you have an example of properly registering BatchedNMS_TRT as a plugin?

BatchedNMS_TRT is shipped as a TRT plugin, so you do not need any extra steps to register it.

Also, do you have end to end example of replacing BatchedNMS with a TRT version?

I don't know of an end-to-end example that uses NMS specifically, but we do have an example of GroupNorm (onnx-packnet).

However, I think the code you have right now should be nearly working if you uncomment these 2 lines:

# graph.cleanup()
# graph.toposort()

You can try the resulting ONNX model with trtexec to check, and it should automatically pick up the plugin.

[qraleq]

@pranavm-nvidia Tnx for the fast response. I commented out the changes to the parser and rebuilt onnx-tensorrt. Now when I uncomment the two lines as you proposed, I get the following error when running with trtexec:

[11/04/2020-21:04:04] [W] [TRT] /home/f/Development/onnx-tensorrt/onnx2trt_utils.cpp:220: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
[11/04/2020-21:04:04] [W] [TRT] /home/f/Development/onnx-tensorrt/onnx2trt_utils.cpp:246: One or more weights outside the range of INT32 was clamped
[11/04/2020-21:04:04] [W] [TRT] /home/f/Development/onnx-tensorrt/onnx2trt_utils.cpp:246: One or more weights outside the range of INT32 was clamped
[11/04/2020-21:04:04] [W] [TRT] /home/f/Development/onnx-tensorrt/onnx2trt_utils.cpp:246: One or more weights outside the range of INT32 was clamped
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/ModelImporter.cpp:139: No importer registered for op: BatchedNMS_TRT. Attempting to import as plugin.
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3775: Searching for plugin: BatchedNMS_TRT, plugin_version: 1, plugin_namespace: 
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3792: Successfully created plugin: BatchedNMS_TRT
[11/04/2020-21:04:04] [E] [TRT] onnx_graphsurgeon_node_0: PluginV2Layer must be V2DynamicExt when there are runtime input dimensions.
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/ModelImporter.cpp:139: No importer registered for op: BatchedNMS_TRT. Attempting to import as plugin.
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3775: Searching for plugin: BatchedNMS_TRT, plugin_version: 1, plugin_namespace: 
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3792: Successfully created plugin: BatchedNMS_TRT
[11/04/2020-21:04:04] [E] [TRT] onnx_graphsurgeon_node_0: PluginV2Layer must be V2DynamicExt when there are runtime input dimensions.
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/ModelImporter.cpp:139: No importer registered for op: BatchedNMS_TRT. Attempting to import as plugin.
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3775: Searching for plugin: BatchedNMS_TRT, plugin_version: 1, plugin_namespace: 
[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3792: Successfully created plugin: BatchedNMS_TRT
[11/04/2020-21:04:04] [E] [TRT] onnx_graphsurgeon_node_0: PluginV2Layer must be V2DynamicExt when there are runtime input dimensions.
[11/04/2020-21:04:04] [E] [TRT] onnx_graphsurgeon_node_0: PluginV2Layer must be V2DynamicExt when there are runtime input dimensions.
[11/04/2020-21:04:04] [E] [TRT] onnx_graphsurgeon_node_0: PluginV2Layer must be V2DynamicExt when there are runtime input dimensions.
[11/04/2020-21:04:04] [E] [TRT] onnx_graphsurgeon_node_0: PluginV2Layer must be V2DynamicExt when there are runtime input dimensions.
While parsing node number 351 [Slice]:
ERROR: /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3154 In function importSlice:
[4] Assertion failed: -r <= axis && axis < r
[11/04/2020-21:04:04] [E] Failed to parse onnx file
[11/04/2020-21:04:04] [E] Parsing model failed
[11/04/2020-21:04:04] [E] Engine creation failed
[11/04/2020-21:04:04] [E] Engine set up failed
&&&& FAILED TensorRT.trtexec # /usr/src/tensorrt/bin/trtexec --onnx=model_gs.onnx

Do you have an idea what might be wrong?

[pranavm-nvidia]

@qraleq Looks like your model is using dynamic shapes somewhere:

[11/04/2020-21:04:04] [E] [TRT] onnx_graphsurgeon_node_0: PluginV2Layer must be V2DynamicExt when there are runtime input dimensions.

Seems like the model input dimensions are fixed though, so it must be due to some intermediate layer whose dimensions can't be resolved at build time.
Can you try using the BatchedNMSDynamic_TRT plugin instead?

The good news is that it does seem to be finding the plugin correctly at least:

[11/04/2020-21:04:04] [I] [TRT] /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3792: Successfully created plugin: BatchedNMS_TRT

Also I think you'll probably need to populate the other parameters expected by the plugin (see here) or you might run into issues during inference.

[qraleq]

@pranavm-nvidia Yeah, I also figured that something "dynamic" is still leftover in the graph, but can't figure out what.

I tried to use BatchedNMSDynamic_TRT, but I get error saying:

[11/04/2020-21:32:10] [I] [TRT] /home/f/Development/onnx-tensorrt/ModelImporter.cpp:139: No importer registered for op: BatchedNMSDynamic_TRT. Attempting to import as plugin.
[11/04/2020-21:32:10] [I] [TRT] /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3775: Searching for plugin: BatchedNMSDynamic_TRT, plugin_version: 1, plugin_namespace: 
[11/04/2020-21:32:10] [E] [TRT] INVALID_ARGUMENT: getPluginCreator could not find plugin BatchedNMSDynamic_TRT version 1
While parsing node number 834 [BatchedNMSDynamic_TRT]:
ERROR: /home/f/Development/onnx-tensorrt/builtin_op_importers.cpp:3777 In function importFallbackPluginImporter:
[8] Assertion failed: creator && "Plugin not found, are the plugin name, version, and namespace correct?"

I'm using TensorRT 7.1.3.0-1+cuda10.2 on Jetson AGX. How can I make sure that dynamic NMS gets recognized by the parser?

[pranavm-nvidia]

@qraleq The dynamic version was added in 7.2, so that would explain why it doesn't recognize it.

Would you be able to share the verbose trtexec logs (add --verbose) and/or your ONNX model?

[qraleq]

@pranavm-nvidia Please find the log attached.
log.txt

Can I send you the model directly somehow?

[pranavm-nvidia]

@qraleq You can email me the model/link to the model at pranavm@nvidia.com

Seems like it's this Resize that's introducing dynamic shapes:

ModelImporter.cpp:107: Parsing node: Resize__1007 [Resize]
ModelImporter.cpp:129: Resize__1007 [Resize] inputs: [Transpose__1002:0 -> (1, 3, 1280, 720)], [roi__998 -> ()], [roi__998 -> ()], [Concat__1006:0 -> (4)], 
ModelImporter.cpp:183: Resize__1007 [Resize] outputs: [Resize__1007:0 -> (-1, -1, -1, -1)], 

Note that the output of the previous layer has a fixed shape: Transpose__1002:0 -> (1, 3, 1280, 720)

Is it possible to express the resize using a constant sizes or scales input instead of rois? Then TRT would be able to compute the output shapes at build time

[mk-nvidia]

@qraleq Please look at @pranavm-nvidia 's comment above to see if it fixes your problem.