search

WongKinYiu / yolor

implementation of paper - You Only Learn One Representation: Unified Network for Multiple Tasks (https://arxiv.org/abs/2105.04206)
GNU General Public License v3.0
1.99k stars 518 forks source link

Issue with onnx output shape #207

Closed jnskplr closed 2 years ago

jnskplr commented 2 years ago

Hi,

I'm on paper branch and trained yolor-p6 on custom dataset. Now I exported it to onnx model.

Shape of my onnx output layer is batch, 3, 96, 160, nc. That differs from shape of forward() function return when doing detect.py.

When I set model.model[-1].export = False in export.py, I will get the correct output layer in my onnx model, but that model does not work for any other inference engine.

Setting model.model[-1].export True or False delivers two different onnx models, why is that the case? Any suggestion of what this problem could solve or how to get the correct output layer in onnx model?

Thanks.

jnskplr commented 2 years ago

Followed steps described in this issue from yolov5: convert yolov5 to openvino

Getting correct output layer now and onnx model works.