From b1a12b49b71a24d3236895ab0fa74638ab63d0da Mon Sep 17 00:00:00 2001 From: Yufeng Li Date: Wed, 17 Feb 2021 19:55:19 -0800 Subject: [PATCH] Avoid removing constant weight that is graph output (#6735) --- .../E2E_example_model/image_classification/cpu/run.py | 5 +++-- onnxruntime/python/tools/quantization/onnx_model.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/onnxruntime/python/tools/quantization/E2E_example_model/image_classification/cpu/run.py b/onnxruntime/python/tools/quantization/E2E_example_model/image_classification/cpu/run.py index 5212967a8d..100057cf95 100644 --- a/onnxruntime/python/tools/quantization/E2E_example_model/image_classification/cpu/run.py +++ b/onnxruntime/python/tools/quantization/E2E_example_model/image_classification/cpu/run.py @@ -12,7 +12,7 @@ from PIL import Image import onnx import onnxruntime from onnx import helper, TensorProto, numpy_helper -from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantFormat +from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantFormat, QuantType class ResNet50DataReader(CalibrationDataReader): @@ -107,7 +107,8 @@ def main(): output_model_path, dr, quant_format=args.quant_format, - per_channel=args.per_channel) + per_channel=args.per_channel, + weight_type=QuantType.QInt8) print('Calibrated and quantized model saved.') print('benchmarking fp32 model...') diff --git a/onnxruntime/python/tools/quantization/onnx_model.py b/onnxruntime/python/tools/quantization/onnx_model.py index 6ba2e99a99..22af9071aa 100644 --- a/onnxruntime/python/tools/quantization/onnx_model.py +++ b/onnxruntime/python/tools/quantization/onnx_model.py @@ -232,14 +232,15 @@ class ONNXModel: unused_nodes = [] nodes = self.nodes() for node in nodes: - if node.op_type == "Constant" and node.output[0] not in input_name_to_nodes: + if node.op_type == "Constant" and not self.is_graph_output( + node.output[0]) and node.output[0] not in input_name_to_nodes: unused_nodes.append(node) self.remove_nodes(unused_nodes) ununsed_weights = [] for w in self.initializer(): - if w.name not in input_name_to_nodes: + if w.name not in input_name_to_nodes and not self.is_graph_output(w.name): ununsed_weights.append(w) # Remove from graph.input for graph_input in self.graph().input: