onnxruntime shape mismatch during quantization of yolov8 models

[Jamil]

Describe the issue

When trying to quantize a Yolov8 model (exported with yolo export model=yolov8x.pt format=onnx) with onnxruntime, I get the following error:

$ python quantize.py yolov8x.onnx
Model changed? False
Model to quantize: ./yolov8x.onnx
Exclude nodes:
[]
WARNING:root:Please consider to run pre-processing before quantization. Refer to example: https://github.com/microsoft/onnxruntime-inference-examples/blob/main/quantization/image_classification/cpu/ReadMe.md

WARNING:root:Please consider pre-processing before quantization. See https://github.com/microsoft/onnxruntime-inference-examples/blob/main/quantization/image_classification/cpu/ReadMe.md
Finished quantization. Validating...
(1, 3, 640, 640)
2024-06-14 10:10:32.328524183 [W:onnxruntime:, execution_frame.cc:660 AllocateMLValueTensorPreAllocateBuffer] Shape mismatch attempting to re-use buffer. {1,40,40,640} != {1,39,39,642}. Validate usage of dim_value (values should be > 0) and dim_param (all values with the same string should equate to the same size) in shapes in the model.
2024-06-14 10:10:32.328606668 [E:onnxruntime:, sequential_executor.cc:516 ExecuteKernel] Non-zero status code returned while running QLinearConcat node. Name:'/model.11/Concat' Status Message: concat.cc:154 PrepareForCompute Non concat axis dimensions must match: Axis 1 has mismatched dimensions of 40 and 39
Traceback (most recent call last):
  File "quantize.py", line 124, in <module>
    quant_outputs = test_model(output_model_path, input_data)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "quantize.py", line 32, in test_model
    outputs = session.run(None, input_data)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/…/anaconda3/envs/yolo/lib/python3.11/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 220, in run
    return self._sess.run(output_names, input_feed, run_options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running QLinearConcat node. Name:'/model.11/Concat' Status Message: concat.cc:154 PrepareForCompute Non concat axis dimensions must match: Axis 1 has mismatched dimensions of 40 and 39

To reproduce

• Export a YOLO model using yolo export model=yolov8x.pt format=onnx
• Quantize the ONNX file as follows:

    qnn_config = get_qnn_qdq_config(model_to_quantize,
                                    data_reader,
                                    activation_type=QuantType.QUInt8,
                                    weight_type=QuantType.QUInt8,
                                    per_channel=True,
                                    activation_symmetric=True,
                                    weight_symmetric=True)

    output_model_path = os.path.join(models_directory, 'model.qdq.onnx')
    quantize(model_to_quantize, output_model_path, qnn_config)

Urgency

This is blocking for the project I'm working on, and seems like a regression in onnxruntime functionality.

Platform

Linux

OS Version

Ubuntu 22.04

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.18.0

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Other / Unknown

Execution Provider Library Version

QNN