search

TrojanXu / yolov5-tensorrt

A tensorrt implementation of yolov5: https://github.com/ultralytics/yolov5
Apache License 2.0
191 stars 46 forks source link

Differing Output Shape After Onnx Export/Simplification #24

Closed seanavery closed 3 years ago

seanavery commented 3 years ago

I am unable to build the TensorRT engine from the onnx export in the yolov5 repo. It looks like I have a different output shape network.get_output(0).

The output shape I am seeing is (1, 3, 80, 80, 85), while in your code it is expecting a three dimensions (bs, boxes, _). https://github.com/TrojanXu/yolov5-tensorrt/blob/master/main.py#L184

The onnx simplification step runs just fine, but wondering how to get to the output shape required by the build_engine function.

TrojanXu commented 3 years ago

Did you use main.py to export onnx model? It's different from the onnx_export script in yolov5 repo.

seanavery commented 3 years ago

Thanks @TrojanXu