mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Avoid removing constant weight that is graph output (#6735)
This commit is contained in:
parent
ea3aee4d5f
commit
b1a12b49b7
2 changed files with 6 additions and 4 deletions
|
|
@ -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...')
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue