open-mmlab / mmdetection

OpenMMLab Detection Toolbox and Benchmark
https://mmdetection.readthedocs.io
Apache License 2.0
29.61k stars 9.47k forks source link

RTMDET How to get "Grided concatenated outputs" without NMS #11240

Open MaximeDebarbat opened 11 months ago

MaximeDebarbat commented 11 months ago

Hi, I am testing and trying to export RtmDet with an output containing the detections, rescaled bounding boxes and label scores as yolo models would which would provide a final tensor of dimension [batch_size, detection_dimension, detection_size]. Where detection_dimension depends on the input size and detection_size would thus correspond to the concatenation of {bounding_boxes, score, label scores} for instance {4, 1, 80} in the case of a coco trained model.

The model has to be TensorRT friendly, thus operation such as NonZero can not be used.

I have tried exports as per defined in the doc, using mmdeploy but BatchedNMS and operation afterwards are not compliant with my requirements.

I also have tried to directly modify base_dense_head.py which is being used by the model unfortunately, it is batch dependent due to this loop, shall I still create my own head or is there anyone who knows how to do it seamlessly ?

613B commented 11 months ago

Adding code below to deploy_config.py, I could get the onnx model that has the networks prior to BatchedNMS. Also, by using trtexec, I could get TensorRT model from this onnx model. After using this technique, you should be able to implement the post-processing yourself and get the model you want.

I don't know if this method will help you.

partition_config = dict(
    type='tensorrt',
    apply_marks=True, 
    partition_cfg=[
        dict(
            save_file='model.onnx', 
            start=['detector_forward:input'],  # [mark_name:input, ...]
            end=['multiclass_nms:input'],  # [mark_name:output, ...]
            output_names=['boxes', 'scores']) # output names
    ])