Tianxiaomo / pytorch-YOLOv4

PyTorch ,ONNX and TensorRT implementation of YOLOv4
Apache License 2.0
4.46k stars 1.49k forks source link

Darknet => ONNX conversion unable to run with dynamic input shapes #502

Open devaale opened 2 years ago

devaale commented 2 years ago

Hello, I've seen other people deal with similar issue, but unable to find the solution to it. I've converted Darknet model => ONNX using demo_darknet2onnx.py file. My export function looks like this:

model = Darknet(cfgfile)
model.print_network()
model.load_weights(weightfile)

input_names = ["input"]
output_names = ['boxes', 'confs']

x = torch.randn((1, 3, model.height, model.width), requires_grad=True)
dynamic_axes =
{
            "input": {0: "batch_size", 2: "height", 3: "width"},
            "boxes": {0: "batch_size", 1: "bounding_boxes"},
            "confs": {0: "batch_size", 1: "scores"}
}

torch.onnx.export(model,
                  x,
                  'model.onnx',
                  export_params=True,
                  opset_version=11,
                  do_constant_folding=True,
                  input_names=input_names,
                  output_names=output_names,
                  dynamic_axes=dynamic_axes)

And I end up with such model.onnx file:

image

However inference only works with model (height, width) - in this case image with shape=1x3x2976x64. Anything else (ex.: shape=1x3x992x64) results in:

outputs = session.run(['boxes', 'confs'], {'input': img_in}) return self._sess.run(output_names, input_feed, run_options) onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Add node. Name:'Add_1514' Status Message: D:\a\_work\1\s\onnxruntime\core/providers/cpu/math/element_wise_ops.h:505 onnxruntime::BroadcastIterator::Append axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 2 by 93 I understand that this part of code fails as stated in error:

        try:
            return self._sess.run(output_names, input_feed, run_options)   #<------- THIS ONE
        except C.EPFail as err:
            if self._enable_fallback:
                print("EP Error: {} using {}".format(str(err), self._providers))
                print("Falling back to {} and retrying.".format(self._fallback_providers))
                self.set_providers(self._fallback_providers)
                # Fallback only once.
                self.disable_fallback()
                return self._sess.run(output_names, input_feed, run_options)

And it fails at node 'Add_1514' which looks like this:

image

I run inference:

session = onnxruntime.InferenceSession(onnx_model_path)
outputs = session.run(None, {'input': img_in})

where img_in is: image

Any help of how I could tackle this problem would be appreciated. In the end I want to convert ONNX to TRT and run inference on triton server, but I need dynamic input shapes so I could pass different size images.

frenky-strasak commented 2 years ago

I suggest you to use "BatchedNMSDynamic_TRT" plugin. It is desribed here: https://github.com/NVIDIA-AI-IOT/yolov4_deepstream/blob/master/tensorrt_yolov4/README.md