Skip to content

Commit

Permalink
Revert "Make keepdims to its default value when adding ReduceMin/Redu…
Browse files Browse the repository at this point in the history
…ceMax for quantization calibration (#6788)" (#6825)

This reverts commit 9b3171e.
  • Loading branch information
oliviajain authored Feb 27, 2021
1 parent 4bbea91 commit 29b30bb
Showing 1 changed file with 10 additions and 14 deletions.
24 changes: 10 additions & 14 deletions onnxruntime/python/tools/quantization/calibrate.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
import numpy as np
import onnx
import onnxruntime
from onnx import helper, TensorProto, ModelProto, shape_inference
from onnx import helper, TensorProto, ModelProto
from onnx import onnx_pb as onnx_proto
from six import string_types
from enum import Enum
Expand Down Expand Up @@ -52,9 +52,6 @@ def __init__(self, model, op_types_to_calibrate=[], augmented_model_path='augmen
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

Expand Down Expand Up @@ -163,27 +160,26 @@ def augment_graph(self):

added_nodes = []
added_outputs = []
tensors, value_infos = self.select_tensors_to_calibrate(model)
tensors, _ = 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)
reduce_min_node = onnx.helper.make_node('ReduceMin', [tensor], [tensor + '_ReduceMin'],
reduce_min_name,
keepdims=0)

added_nodes.append(reduce_min_node)
added_outputs.append(helper.make_tensor_value_info(reduce_min_node.output[0], TensorProto.FLOAT, shape))
added_outputs.append(helper.make_tensor_value_info(reduce_min_node.output[0], TensorProto.FLOAT, ()))

# Adding ReduceMax nodes
reduce_max_name = tensor + '_ReduceMax'
reduce_max_node = onnx.helper.make_node('ReduceMax', [tensor], [tensor + '_ReduceMax'], reduce_max_name)
reduce_max_node = onnx.helper.make_node('ReduceMax', [tensor], [tensor + '_ReduceMax'],
reduce_max_name,
keepdims=0)

added_nodes.append(reduce_max_node)
added_outputs.append(helper.make_tensor_value_info(reduce_max_node.output[0], TensorProto.FLOAT, shape))
added_outputs.append(helper.make_tensor_value_info(reduce_max_node.output[0], TensorProto.FLOAT, ()))

model.graph.node.extend(added_nodes)
model.graph.output.extend(added_outputs)
Expand Down

0 comments on commit 29b30bb

Please sign in to comment.