[Quantization] Apply workaround for crash when using histogram-based calibrators (#21972)

### Description
- Applies a workaround that prevents the histogram-based calibrators
(percentile, entropy, distribution) from crashing. The workaround
involves copying inference outputs that come directly from model inputs.
A description of the bug is here:
https://github.com/microsoft/onnxruntime/issues/21922. **This PR does
not fix the root bug, but instead provides a workaround to _unblock_
users using histogram-based calibration.**
- Adds a unit test that runs all histogram-based calibrators to help
catch future regressions. We didn't have unit tests that ran these
calibration methods.

### Motivation and Context
Trying to quantize a model with the percentile, entropy, or distribution
calibration methods raises an exception:
```shell
  File "/.../site-packages/onnxruntime/quantization/quantize.py", line 691, in quantize
    quantize_static(
  File "/.../site-packages/onnxruntime/quantization/quantize.py", line 525, in quantize_static
    calibrator.collect_data(calibration_data_reader)
  File "/.../site-packages/onnxruntime/quantization/calibrate.py", line 571, in collect_data
    self.collector.collect(clean_merged_dict)
  File "/.../site-packages/onnxruntime/quantization/calibrate.py", line 746, in collect
    return self.collect_value(name_to_arr)
  File "/.../site-packages/onnxruntime/quantization/calibrate.py", line 836, in collect_value
    hist, hist_edges = np.histogram(data_arr, self.num_bins, range=(-threshold, threshold))
  File "<__array_function__ internals>", line 180, in histogram
  File ".../site-packages/numpy/lib/histograms.py", line 793, in histogram
    bin_edges, uniform_bins = _get_bin_edges(a, bins, range, weights)
  File "/.../site-packages/numpy/lib/histograms.py", line 426, in _get_bin_edges
    first_edge, last_edge = _get_outer_edges(a, range)
  File "/.../site-packages/numpy/lib/histograms.py", line 315, in _get_outer_edges
    raise ValueError(
ValueError: supplied range of [nan, nan] is not finite
```

The calibrators create an augmented model with all tensors (including
model inputs) set as model outputs. The data for outputs that are also
model inputs is corrupted as described in
https://github.com/microsoft/onnxruntime/issues/21922. The corrupted
data sometimes contains `NaN` values that cause numpy's histogram
utilities to raise an exception.
This commit is contained in:
Adrian Lizarraga 2024-09-09 12:05:41 -07:00 committed by GitHub
parent 2cdc05f189
commit c7ae9b977a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 59 additions and 11 deletions

View file

@ -565,16 +565,29 @@ class HistogramCalibrater(CalibraterBase):
"""
Entropy Calibrator collects operators' tensors as well as generates tensor histogram for each operator.
"""
input_names_set = {node_arg.name for node_arg in self.infer_session.get_inputs()}
output_names = [node_arg.name for node_arg in self.infer_session.get_outputs()]
while True:
inputs = data_reader.get_next()
if not inputs:
break
self.intermediate_outputs.append(self.infer_session.run(None, inputs))
outputs = self.infer_session.run(None, inputs)
# Copy np.ndarray only for graph outputs that are also graph inputs to workaround bug:
# https://github.com/microsoft/onnxruntime/issues/21922
fixed_outputs = []
for output_index, output in enumerate(outputs):
if output_names[output_index] in input_names_set:
fixed_outputs.append(copy.copy(output))
else:
fixed_outputs.append(output)
self.intermediate_outputs.append(fixed_outputs)
if len(self.intermediate_outputs) == 0:
raise ValueError("No data is collected.")
output_names = [self.infer_session.get_outputs()[i].name for i in range(len(self.intermediate_outputs[0]))]
output_dicts_list = [
dict(zip(output_names, intermediate_output)) for intermediate_output in self.intermediate_outputs
]

View file

@ -14,7 +14,7 @@ import onnx
from onnx import TensorProto, helper, numpy_helper
import onnxruntime
from onnxruntime.quantization.calibrate import CalibrationDataReader, create_calibrator
from onnxruntime.quantization.calibrate import CalibrationDataReader, CalibrationMethod, create_calibrator
def generate_input_initializer(tensor_shape, tensor_dtype, input_name):
@ -275,7 +275,7 @@ class TestCalibrateMinMaxCalibrator(unittest.TestCase):
for output in added_outputs:
self.assertTrue(output in augmented_model_outputs)
def construct_test_compute_data_model(self, test_model_path, opset_version=13):
def construct_test_compute_data_model(self, test_model_path, opset_version=13, augmented=True):
# (input)
# |
# Relu
@ -290,12 +290,19 @@ class TestCalibrateMinMaxCalibrator(unittest.TestCase):
# |
# (X6)
input = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 1, 3])
x1_output = helper.make_tensor_value_info("X1", TensorProto.FLOAT, [1, 3, 1, 3])
x2_output = helper.make_tensor_value_info("X2", TensorProto.FLOAT, [1, 3, 1, 3])
x3_output = helper.make_tensor_value_info("X3", TensorProto.FLOAT, [1, 3, 1, 3])
x4_output = helper.make_tensor_value_info("X4", TensorProto.FLOAT, [1, 3, 1, 3])
x5_output = helper.make_tensor_value_info("X5", TensorProto.FLOAT, [1, 3, 1, 3])
x6_output = helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3])
graph_outputs = None
if augmented:
graph_outputs = [
helper.make_tensor_value_info("X1", TensorProto.FLOAT, [1, 3, 1, 3]),
helper.make_tensor_value_info("X2", TensorProto.FLOAT, [1, 3, 1, 3]),
helper.make_tensor_value_info("X3", TensorProto.FLOAT, [1, 3, 1, 3]),
helper.make_tensor_value_info("X4", TensorProto.FLOAT, [1, 3, 1, 3]),
helper.make_tensor_value_info("X5", TensorProto.FLOAT, [1, 3, 1, 3]),
helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3]),
]
else:
graph_outputs = [helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3])]
w1 = generate_input_initializer([3, 3, 1, 1], np.float32, "W1")
b1 = generate_input_initializer([3], np.float32, "B1")
w3 = generate_input_initializer([3, 3, 1, 1], np.float32, "W3")
@ -312,7 +319,7 @@ class TestCalibrateMinMaxCalibrator(unittest.TestCase):
[relu_node_1, conv_node_1, relu_node_2, conv_node_2, conv_node_3, add_node],
"test_graph_4",
[input],
[x1_output, x2_output, x3_output, x4_output, x5_output, x6_output],
graph_outputs,
)
graph.initializer.add().CopyFrom(w1)
graph.initializer.add().CopyFrom(b1)
@ -357,6 +364,34 @@ class TestCalibrateMinMaxCalibrator(unittest.TestCase):
for output_name in output_min_max_dict:
self.assertEqual(output_min_max_dict[output_name], tensors_range[output_name].range_value)
def test_histogram_calibrators_run(self):
"""
Runs all histogram-based calibrators (Percentile, Entropy, Distribution) and checks that they run
and generate the expected number of tensor ranges. Does not check correctness of range values.
"""
# Create test model.
test_model_path = Path(self._tmp_model_dir.name).joinpath("./test_model_4.onnx")
self.construct_test_compute_data_model(test_model_path.as_posix(), augmented=False)
# Count the number of tensors in the model.
model = onnx.load_model(test_model_path)
model = onnx.shape_inference.infer_shapes(model)
num_tensors = len(model.graph.value_info) + len(model.graph.input) + len(model.graph.output)
# Run all histogram calibration methods.
data_reader = TestDataReader()
calibration_methods = [CalibrationMethod.Percentile, CalibrationMethod.Entropy, CalibrationMethod.Distribution]
for calibration_method in calibration_methods:
with self.subTest(calibration_method=calibration_method):
data_reader.rewind()
augmented_model_path = Path(self._tmp_model_dir.name).joinpath(f"augmented_{calibration_method}.onnx")
calibrator = create_calibrator(
test_model_path, calibrate_method=calibration_method, augmented_model_path=augmented_model_path
)
calibrator.collect_data(data_reader)
tensors_range = calibrator.compute_data()
self.assertEqual(len(tensors_range.items()), num_tensors) # A range for every tensor in the graph.
def test_augment_graph_with_zero_value_dimension(self):
"""TEST_CONFIG_5"""
# Conv