[Training] quantize_static Histogram based calibration methods expect fixed shape outputs

[kmn1024]

Describe the issue

When using quantize_static on models that produce variable shaped outputs, using calibrate methods that need to collect a histogram first will run this line and fail:

onnxruntime/onnxruntime/python/tools/quantization/calibrate.py

Line 719 in d673e39

data_arr = np.asarray(data_arr) # noqa: PLW2901

For example, I had a model which outputs (Batch Size, Sequence Length, Embedding dims)

data_arr = np.asarray(data_arr) # noqa: PLW2901
ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 2 dimensions. The detected shape was (200, 1) + inhomogeneous part.

To reproduce

Run this on a model that produces variable sized outputs:

quantize_static('./onnx/xxx.prep.onnx',
                                   './onnx/quantized/xxx.prep.static-dist.onnx',
                                   calibration_data_reader=cal_data,
                                   calibrate_method=CalibrationMethod.Distribution)

Urgency

quantize_static still works with CalibrationMethod.MinMax, and quantize_dynamic also works. I just want to see how well the other calibration method works for my model.

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

22.04

PyTorch Version

2.1.1

Execution Provider

Default CPU

Execution Provider Library Version

No response

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.

[volker48]

Just hit this issue as well trying to use CalibrationMethod.Entropy with dynamic_axes for the batch size.