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

Make keepdims to its default value when adding ReduceMin/ReduceMax for quantization calibration (#6788)

Browse source 
* Make keepdims to its default value when adding ReduceMin/ReduceMax

* Fix bug for adding ReduceMin/ReduceMax with keepdims=1
This commit is contained in:
Chi Lo 2021-02-25 09:47:59 -08:00 committed by GitHub
parent db05d53b94
commit 9b3171e95c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -10,7 +10,7 @@ import os
import numpy as np
import onnx
import onnxruntime
from onnx import helper, TensorProto, ModelProto
from onnx import helper, TensorProto, ModelProto, shape_inference
from onnx import onnx_pb as onnx_proto
from six import string_types
from enum import Enum
@ -52,6 +52,9 @@ class CalibraterBase:
else:
raise ValueError('model should be either model path or onnx.ModelProto.')
# Apply shape inference on the model
self.model = onnx.shape_inference.infer_shapes(self.model)
self.op_types_to_calibrate = op_types_to_calibrate
self.augmented_model_path = augmented_model_path
@ -160,26 +163,27 @@ class MinMaxCalibrater(CalibraterBase):
added_nodes = []
added_outputs = []
tensors, _ = self.select_tensors_to_calibrate(model)
tensors, value_infos = self.select_tensors_to_calibrate(model)
for tensor in tensors:
# Get tensor's shape
dim = len(value_infos[tensor].type.tensor_type.shape.dim)
shape = (1,) if dim == 1 else list(1 for i in range(dim))
# Adding ReduceMin nodes
reduce_min_name = tensor + '_ReduceMin'
reduce_min_node = onnx.helper.make_node('ReduceMin', [tensor], [tensor + '_ReduceMin'],
reduce_min_name,
keepdims=0)
reduce_min_node = onnx.helper.make_node('ReduceMin', [tensor], [tensor + '_ReduceMin'], reduce_min_name)
added_nodes.append(reduce_min_node)
added_outputs.append(helper.make_tensor_value_info(reduce_min_node.output[0], TensorProto.FLOAT, ()))
added_outputs.append(helper.make_tensor_value_info(reduce_min_node.output[0], TensorProto.FLOAT, shape))
# Adding ReduceMax nodes
reduce_max_name = tensor + '_ReduceMax'
reduce_max_node = onnx.helper.make_node('ReduceMax', [tensor], [tensor + '_ReduceMax'],
reduce_max_name,
keepdims=0)
reduce_max_node = onnx.helper.make_node('ReduceMax', [tensor], [tensor + '_ReduceMax'], reduce_max_name)
added_nodes.append(reduce_max_node)
added_outputs.append(helper.make_tensor_value_info(reduce_max_node.output[0], TensorProto.FLOAT, ()))
added_outputs.append(helper.make_tensor_value_info(reduce_max_node.output[0], TensorProto.FLOAT, shape))
model.graph.node.extend(added_nodes)
model.graph.output.extend(added_outputs)