diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index 1a1bafea4f..eeccdc6868 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -58,6 +58,7 @@ class ONNXCalibrater: model = onnx.shape_inference.infer_shapes(model) value_infos = {vi.name: vi for vi in model.graph.value_info} value_infos.update({ot.name: ot for ot in model.graph.output}) + value_infos.update({it.name: it for it in model.graph.input}) added_nodes = [] added_outputs = [] @@ -264,4 +265,4 @@ def calibrate(model_path, quantization_params_dict = calibrater.calculate_quantization_params(dict_for_quantization) print("Calibrated,quantized parameters calculated and returned.") - return quantization_params_dict \ No newline at end of file + return quantization_params_dict diff --git a/onnxruntime/python/tools/quantization/onnx_quantizer.py b/onnxruntime/python/tools/quantization/onnx_quantizer.py index abe10af477..661621cc6b 100644 --- a/onnxruntime/python/tools/quantization/onnx_quantizer.py +++ b/onnxruntime/python/tools/quantization/onnx_quantizer.py @@ -806,7 +806,7 @@ class ONNXQuantizer: # Quantize the input initializer = find_by_name(node_input, self.model.initializer()) if initializer is not None: - weight = self._get_quantized_weight(initializer, self.weight_qType) + weight = self._get_quantized_weight(initializer, self.weight_qType if initializer_use_weight_qType else self.input_qType) # Update graph self._update_weight(weight)