saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Calibrate float tensor only (#5704)

Browse source
This commit is contained in:
Yufeng Li 2020-11-04 23:55:48 -08:00 committed by GitHub
parent 2127a229d7
commit 5c4543e194
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
onnxruntime/python/tools/quantization/calibrate.py
View file

@ -12,6 +12,7 @@ import onnxruntime
from onnx import helper, TensorProto
import abc
import itertools
class CalibrationDataReader(metaclass=abc.ABCMeta):
@ -54,6 +55,9 @@ class ONNXCalibrater:
'''
model = onnx.load(self.model_path)
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})
added_nodes = []
added_outputs = []
@ -63,12 +67,12 @@ class ONNXCalibrater:
should_be_calibrate = ((node.op_type in self.calibrate_op_types) and
(node.name not in self.black_nodes)) or (node.name in self.white_nodes)
if should_be_calibrate:
tensors_to_calibrate.update(node.input)
tensors_to_calibrate.update(node.output)
for tensor in tensors_to_calibrate:
if tensor in model.graph.initializer:
tensors_to_calibrate.remove(tensor)
for tensor_name in itertools.chain(node.input, node.output):
if tensor_name in value_infos.keys():
vi = value_infos[tensor_name]
if vi.type.HasField('tensor_type') and vi.type.tensor_type.elem_type == TensorProto.FLOAT and (
tensor_name not in model.graph.initializer):
tensors_to_calibrate.add(tensor_name)
for tensor in tensors_to_calibrate:
# Adding ReduceMin nodes