Avoid removing constant weight that is graph output (#6735)

This commit is contained in:
Yufeng Li 2021-02-17 19:55:19 -08:00 committed by GitHub
parent ea3aee4d5f
commit b1a12b49b7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 4 deletions

View file

@ -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...')

View file

@ -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: