Tianxiaomo / pytorch-YOLOv4

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

If you are interested in deploying ONNX of darknet model (YOLOv3/YOLOv4/Scaled-YOLOv4) #484

Open james77777778 opened 2 years ago

james77777778 commented 2 years ago

I try to refactor the ONNX conversion codes in this repo and only put the focus on the deployment (not training!). I think the AlexeyAB/darknet framework is easy enough and more stable than any other yolo-like project.

https://github.com/james77777778/darknet-onnx

Also, it should be easy to convert to tensorrt by trtexec if you got .onnx format. I will try to implement tensorrt conversion in the future.

devaale commented 2 years ago

@james77777778 Have you tried converting Darket to ONNX with dynamic input / output shapes ? (Like width, height)

I've converted darknet to onnx with dynamic_axes parameter:

x = torch.randn((1, 3, model.height, model.width), requires_grad=True)
### model.height = 2976, model.width = 64
input_names = ["input"]
output_names = ['boxes', 'confs']

dynamic_axes = {
            "input": {0:"batch_size", 2: "height", 3: "width"},
            "boxes": {0:"batch_size"},
            "confs": {0:"batch_size"}            
}
torch.onnx.export(model,
                          x,
                          onnx_file_name,
                          export_params=True,
                          opset_version=11,
                          do_constant_folding=True,
                          input_names=input_names,
                          output_names=output_names,
                          dynamic_axes=dynamic_axes)

And end up with: image However when I later convert to TRT with the command:

trtexec --explicitBatch --onnx=/yolov4-triton-tensorrt/yolov4_-1_3_2976_64_dynamic.onnx --minShapes=input:1x3x416x64 --optShapes=input:1x3x2976x64 --maxShapes=input:1x3x9984x64 --shapes=input:1x3x2976x64 --saveEngine=/yolov4-triton-tensorrt/yolov4_dynamic.engine --workspace=3000 --fp16

Only inputs with 1x3x2976x64 shape work... Any ideas why could that be ? - Maybe I'm making obvious mistake that is clearly visible ?