fix output shape of ReduceMin/ReduceMax in calibration tool (#11647)

This commit is contained in:
Yufeng Li 2022-05-31 14:26:08 -07:00 committed by GitHub
parent 004560f1fe
commit f437945926
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 150 additions and 156 deletions

View file

@ -7,25 +7,17 @@
# --------------------------------------------------------------------------
import abc
import itertools
import uuid
from enum import Enum
from pathlib import Path
import numpy as np
import onnx
from onnx import ModelProto, TensorProto, helper
from onnx import onnx_pb as onnx_proto
from onnx import ModelProto, TensorProto, helper, numpy_helper
import onnxruntime
from .quant_utils import (
QuantType,
apply_plot,
clone_model_with_shape_infer,
load_model,
model_has_infer_metadata,
smooth_distribution,
)
from .registry import QLinearOpsRegistry
from .quant_utils import apply_plot, clone_model_with_shape_infer, load_model, smooth_distribution
class CalibrationMethod(Enum):
@ -98,9 +90,7 @@ class CalibraterBase:
sess_options = onnxruntime.SessionOptions()
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
self.infer_session = onnxruntime.InferenceSession(
self.augmented_model_path,
sess_options=sess_options,
providers=self.execution_providers,
self.augmented_model_path, sess_options=sess_options, providers=self.execution_providers,
)
def select_tensors_to_calibrate(self, model):
@ -180,11 +170,7 @@ class MinMaxCalibrater(CalibraterBase):
:param averaging_constant: constant smoothing factor to use when computing the moving average.
"""
super(MinMaxCalibrater, self).__init__(
model,
op_types_to_calibrate,
augmented_model_path,
symmetric,
use_external_data_format,
model, op_types_to_calibrate, augmented_model_path, symmetric, use_external_data_format,
)
self.intermediate_outputs = []
self.calibrate_tensors_range = None
@ -203,56 +189,39 @@ class MinMaxCalibrater(CalibraterBase):
"""
model = clone_model_with_shape_infer(self.model)
added_nodes = []
added_outputs = []
tensors, value_infos = self.select_tensors_to_calibrate(model)
for tensor in tensors:
tensors, _ = self.select_tensors_to_calibrate(model)
reshape_shape_name = str(uuid.uuid4())
reshape_shape = numpy_helper.from_array(np.array([1], dtype=np.int64), reshape_shape_name)
model.graph.initializer.append(reshape_shape)
def add_reduce_min_max(tensor_name, reduce_op_name):
# When doing ReduceMax/ReduceMin, ORT can't reduce on dim with value of 0 if 'keepdims' is false.
# To make the code simple, we always let keepdims to be 1.
keepdims = 1
# dim could be:
# [dim_param: "batch_size", dim_value: 256, dim_value: 36, dim_value: 64],
# [dim_value: 0],
# ...
# Please see the definition of TensorShapeProto https://github.com/onnx/onnx/blob/master/onnx/onnx.proto#L651
dim = value_infos[tensor].type.tensor_type.shape.dim
shape = (1,) if len(dim) == 1 else tuple(1 for i in range(len(dim)))
# Adding ReduceMin nodes
reduce_min_name = tensor + "_ReduceMin"
reduce_min_node = onnx.helper.make_node(
"ReduceMin",
[tensor],
[tensor + "_ReduceMin"],
reduce_min_name,
keepdims=keepdims,
# Adding ReduceMin/ReduceMax nodes: ReduceMin/ReduceMax -> Reshape-> (output)
reduce_output = tensor_name + "_" + reduce_op_name
intermediate_output = reduce_output + "_Reshape"
reduce_node = onnx.helper.make_node(
reduce_op_name, [tensor_name], [intermediate_output], keepdims=keepdims, name=reduce_output
)
added_nodes.append(reduce_min_node)
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=keepdims,
reshape_node = onnx.helper.make_node(
"Reshape",
inputs=[intermediate_output, reshape_shape_name],
outputs=[reduce_output],
name=intermediate_output,
)
added_nodes.append(reduce_max_node)
added_outputs.append(helper.make_tensor_value_info(reduce_max_node.output[0], TensorProto.FLOAT, shape))
model.graph.node.extend([reduce_node, reshape_node])
model.graph.output.append(helper.make_tensor_value_info(reduce_output, TensorProto.FLOAT, [1]))
for tensor in tensors:
add_reduce_min_max(tensor, "ReduceMin")
add_reduce_min_max(tensor, "ReduceMax")
model.graph.node.extend(added_nodes)
model.graph.output.extend(added_outputs)
onnx.save(
model,
self.augmented_model_path,
save_as_external_data=self.use_external_data_format,
model, self.augmented_model_path, save_as_external_data=self.use_external_data_format,
)
self.augment_model = model
@ -399,9 +368,7 @@ class HistogramCalibrater(CalibraterBase):
model.graph.node.extend(added_nodes)
model.graph.output.extend(added_outputs)
onnx.save(
model,
self.augmented_model_path,
save_as_external_data=self.use_external_data_format,
model, self.augmented_model_path, save_as_external_data=self.use_external_data_format,
)
self.augment_model = model

View file

@ -57,21 +57,16 @@ class TestCalibrate(unittest.TestCase):
# |
# MatMul
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [1, 1, 5, 5])
B = helper.make_tensor_value_info("B", TensorProto.FLOAT, [1, 1, 3, 3])
E = helper.make_tensor_value_info("E", TensorProto.FLOAT, [1, 1, 5, 1])
F = helper.make_tensor_value_info("F", TensorProto.FLOAT, [1, 1, 5, 1])
vi_a = helper.make_tensor_value_info("A", TensorProto.FLOAT, [1, 1, 5, 5])
vi_b = helper.make_tensor_value_info("B", TensorProto.FLOAT, [1, 1, 3, 3])
vi_e = helper.make_tensor_value_info("E", TensorProto.FLOAT, [1, 1, 5, 1])
vi_f = helper.make_tensor_value_info("F", TensorProto.FLOAT, [1, 1, 5, 1])
conv_node = onnx.helper.make_node(
"Conv",
["A", "B"],
["C"],
name="Conv",
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
"Conv", ["A", "B"], ["C"], name="Conv", kernel_shape=[3, 3], pads=[1, 1, 1, 1],
)
clip_node = onnx.helper.make_node("Clip", ["C"], ["D"], name="Clip")
matmul_node = onnx.helper.make_node("MatMul", ["D", "E"], ["F"], name="MatMul")
graph = helper.make_graph([conv_node, clip_node, matmul_node], "test_graph_1", [A, B, E], [F])
graph = helper.make_graph([conv_node, clip_node, matmul_node], "test_graph_1", [vi_a, vi_b, vi_e], [vi_f])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
test_model_path = "./test_model_1.onnx"
@ -86,23 +81,47 @@ class TestCalibrate(unittest.TestCase):
augmented_model_node_names = [node.name for node in augmented_model.graph.node]
augmented_model_outputs = [output.name for output in augmented_model.graph.output]
added_node_names = [
"A_ReduceMin",
"A_ReduceMax",
"B_ReduceMin",
"B_ReduceMax",
"C_ReduceMin",
"C_ReduceMax",
"D_ReduceMin",
"D_ReduceMax",
"E_ReduceMin",
"E_ReduceMax",
"F_ReduceMin",
"F_ReduceMax",
"A_ReduceMin_Reshape",
"A_ReduceMax_Reshape",
"B_ReduceMin_Reshape",
"B_ReduceMax_Reshape",
"C_ReduceMin_Reshape",
"C_ReduceMax_Reshape",
"D_ReduceMin_Reshape",
"D_ReduceMax_Reshape",
"E_ReduceMin_Reshape",
"E_ReduceMax_Reshape",
"F_ReduceMin_Reshape",
"F_ReduceMax_Reshape",
]
added_outputs = [
"A_ReduceMin",
"A_ReduceMax",
"B_ReduceMin",
"B_ReduceMax",
"C_ReduceMin",
"C_ReduceMax",
"D_ReduceMin",
"D_ReduceMax",
"E_ReduceMin",
"E_ReduceMax",
"F_ReduceMin",
"F_ReduceMax",
]
# Original 3 nodes + added ReduceMin/Max nodes
self.assertEqual(len(augmented_model_node_names), 15)
self.assertEqual(len(augmented_model_node_names), 27)
# Original 1 graph output + added outputs * 6
self.assertEqual(len(augmented_model_outputs), 13)
for name in added_node_names:
@ -116,27 +135,17 @@ class TestCalibrate(unittest.TestCase):
# |
# Conv
G = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 1, 5, 5])
H = helper.make_tensor_value_info("H", TensorProto.FLOAT, [1, 1, 3, 3])
J = helper.make_tensor_value_info("J", TensorProto.FLOAT, [1, 1, 3, 3])
K = helper.make_tensor_value_info("K", TensorProto.FLOAT, [1, 1, 5, 5])
vi_g = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 1, 5, 5])
vi_h = helper.make_tensor_value_info("H", TensorProto.FLOAT, [1, 1, 3, 3])
vi_j = helper.make_tensor_value_info("J", TensorProto.FLOAT, [1, 1, 3, 3])
vi_k = helper.make_tensor_value_info("K", TensorProto.FLOAT, [1, 1, 5, 5])
conv_node_1 = onnx.helper.make_node(
"Conv",
["G", "H"],
["I"],
name="Conv1",
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
"Conv", ["G", "H"], ["I"], name="Conv1", kernel_shape=[3, 3], pads=[1, 1, 1, 1],
)
conv_node_2 = onnx.helper.make_node(
"Conv",
["I", "J"],
["K"],
name="Conv2",
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
"Conv", ["I", "J"], ["K"], name="Conv2", kernel_shape=[3, 3], pads=[1, 1, 1, 1],
)
graph = helper.make_graph([conv_node_1, conv_node_2], "test_graph_2", [G, H, J], [K])
graph = helper.make_graph([conv_node_1, conv_node_2], "test_graph_2", [vi_g, vi_h, vi_j], [vi_k])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
test_model_path = "./test_model_2.onnx"
onnx.save(model, test_model_path)
@ -149,9 +158,9 @@ class TestCalibrate(unittest.TestCase):
augmented_model_outputs = [output.name for output in augmented_model.graph.output]
added_node_names = ["I_ReduceMin", "I_ReduceMax", "K_ReduceMin", "K_ReduceMax"]
added_outputs = ["I_ReduceMin", "I_ReduceMax", "K_ReduceMin", "K_ReduceMax"]
# Original 2 nodes + added ReduceMin/Max nodes * 4
self.assertEqual(len(augmented_model_node_names), 12)
# Original 1 graph output + added outputs * 4
# Original 2 nodes + (ReduceMin + Reshape, ReduceMax + Reshape) * 5 tensors
self.assertEqual(len(augmented_model_node_names), 22)
# Original 1 graph output + 5 tensors * 2
self.assertEqual(len(augmented_model_outputs), 11)
for name in added_node_names:
self.assertTrue(name in augmented_model_node_names)
@ -173,21 +182,16 @@ class TestCalibrate(unittest.TestCase):
# |
# (output)
L = helper.make_tensor_value_info("L", TensorProto.FLOAT, [1, 1, 5, 5])
N = helper.make_tensor_value_info("N", TensorProto.FLOAT, [1, 1, 3, 3])
Q = helper.make_tensor_value_info("Q", TensorProto.FLOAT, [1, 1, 5, 5])
vi_l = helper.make_tensor_value_info("L", TensorProto.FLOAT, [1, 1, 5, 5])
vi_n = helper.make_tensor_value_info("N", TensorProto.FLOAT, [1, 1, 3, 3])
vi_q = helper.make_tensor_value_info("Q", TensorProto.FLOAT, [1, 1, 5, 5])
relu_node = onnx.helper.make_node("Relu", ["L"], ["M"], name="Relu")
conv_node = onnx.helper.make_node(
"Conv",
["M", "N"],
["O"],
name="Conv",
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
"Conv", ["M", "N"], ["O"], name="Conv", kernel_shape=[3, 3], pads=[1, 1, 1, 1],
)
clip_node = onnx.helper.make_node("Clip", ["O"], ["P"], name="Clip")
matmul_node = onnx.helper.make_node("MatMul", ["P", "M"], ["Q"], name="MatMul")
graph = helper.make_graph([relu_node, conv_node, clip_node, matmul_node], "test_graph_3", [L, N], [Q])
graph = helper.make_graph([relu_node, conv_node, clip_node, matmul_node], "test_graph_3", [vi_l, vi_n], [vi_q])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
test_model_path = "./test_model_3.onnx"
onnx.save(model, test_model_path)
@ -202,16 +206,30 @@ class TestCalibrate(unittest.TestCase):
added_node_names = [
"M_ReduceMin",
"M_ReduceMax",
"N_ReduceMin",
"N_ReduceMax",
"O_ReduceMin",
"O_ReduceMax",
"P_ReduceMin",
"P_ReduceMax",
"Q_ReduceMin",
"Q_ReduceMax",
"M_ReduceMin_Reshape",
"M_ReduceMax_Reshape",
"N_ReduceMin_Reshape",
"N_ReduceMax_Reshape",
"O_ReduceMin_Reshape",
"O_ReduceMax_Reshape",
"P_ReduceMin_Reshape",
"P_ReduceMax_Reshape",
"Q_ReduceMin_Reshape",
"Q_ReduceMax_Reshape",
]
added_outputs = [
"M_ReduceMin",
"M_ReduceMax",
"N_ReduceMin",
"N_ReduceMax",
"O_ReduceMin",
"O_ReduceMax",
"P_ReduceMin",
@ -219,9 +237,9 @@ class TestCalibrate(unittest.TestCase):
"Q_ReduceMin",
"Q_ReduceMax",
]
# Original 4 nodes + added ReduceMin/Max nodes
self.assertEqual(len(augmented_model_node_names), 14)
# Original 1 graph output + added outputs
# Original 4 nodes + (ReduceMin + Reshape, ReduceMax + Reshape) * 5 tensors
self.assertEqual(len(augmented_model_node_names), 24)
# Original 1 graph output + 5 tensors * 2
self.assertEqual(len(augmented_model_outputs), 11)
for name in added_node_names:
self.assertTrue(name in augmented_model_node_names)
@ -243,18 +261,18 @@ class TestCalibrate(unittest.TestCase):
# |
# (X6)
input = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 1, 3])
X1_output = helper.make_tensor_value_info("X1", TensorProto.FLOAT, [1, 3, 1, 3])
X2_output = helper.make_tensor_value_info("X2", TensorProto.FLOAT, [1, 3, 1, 3])
X3_output = helper.make_tensor_value_info("X3", TensorProto.FLOAT, [1, 3, 1, 3])
X4_output = helper.make_tensor_value_info("X4", TensorProto.FLOAT, [1, 3, 1, 3])
X5_output = helper.make_tensor_value_info("X5", TensorProto.FLOAT, [1, 3, 1, 3])
X6_output = helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3])
W1 = generate_input_initializer([3, 3, 1, 1], np.float32, "W1")
B1 = generate_input_initializer([3], np.float32, "B1")
W3 = generate_input_initializer([3, 3, 1, 1], np.float32, "W3")
B3 = generate_input_initializer([3], np.float32, "B3")
W5 = generate_input_initializer([3, 3, 1, 1], np.float32, "W5")
B5 = generate_input_initializer([3], np.float32, "B5")
x1_output = helper.make_tensor_value_info("X1", TensorProto.FLOAT, [1, 3, 1, 3])
x2_output = helper.make_tensor_value_info("X2", TensorProto.FLOAT, [1, 3, 1, 3])
x3_output = helper.make_tensor_value_info("X3", TensorProto.FLOAT, [1, 3, 1, 3])
x4_output = helper.make_tensor_value_info("X4", TensorProto.FLOAT, [1, 3, 1, 3])
x5_output = helper.make_tensor_value_info("X5", TensorProto.FLOAT, [1, 3, 1, 3])
x6_output = helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3])
w1 = generate_input_initializer([3, 3, 1, 1], np.float32, "W1")
b1 = generate_input_initializer([3], np.float32, "B1")
w3 = generate_input_initializer([3, 3, 1, 1], np.float32, "W3")
b3 = generate_input_initializer([3], np.float32, "B3")
w5 = generate_input_initializer([3, 3, 1, 1], np.float32, "W5")
b5 = generate_input_initializer([3], np.float32, "B5")
relu_node_1 = onnx.helper.make_node("Relu", ["input"], ["X1"], name="Relu1")
conv_node_1 = onnx.helper.make_node("Conv", ["X1", "W1", "B1"], ["X2"], name="Conv1")
relu_node_2 = onnx.helper.make_node("Relu", ["X2"], ["X3"], name="Relu2")
@ -265,14 +283,14 @@ class TestCalibrate(unittest.TestCase):
[relu_node_1, conv_node_1, relu_node_2, conv_node_2, conv_node_3, add_node],
"test_graph_4",
[input],
[X1_output, X2_output, X3_output, X4_output, X5_output, X6_output],
[x1_output, x2_output, x3_output, x4_output, x5_output, x6_output],
)
graph.initializer.add().CopyFrom(W1)
graph.initializer.add().CopyFrom(B1)
graph.initializer.add().CopyFrom(W3)
graph.initializer.add().CopyFrom(B3)
graph.initializer.add().CopyFrom(W5)
graph.initializer.add().CopyFrom(B5)
graph.initializer.add().CopyFrom(w1)
graph.initializer.add().CopyFrom(b1)
graph.initializer.add().CopyFrom(w3)
graph.initializer.add().CopyFrom(b3)
graph.initializer.add().CopyFrom(w5)
graph.initializer.add().CopyFrom(b5)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
onnx.save(model, test_model_path)
@ -289,9 +307,7 @@ class TestCalibrate(unittest.TestCase):
sess_options = onnxruntime.SessionOptions()
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
infer_session = onnxruntime.InferenceSession(
test_model_path,
sess_options=sess_options,
providers=["CPUExecutionProvider"],
test_model_path, sess_options=sess_options, providers=["CPUExecutionProvider"],
)
data_reader.rewind()
rmin = np.array([np.inf, np.inf, np.inf, np.inf, np.inf, np.inf], dtype=np.float32)
@ -318,36 +334,23 @@ class TestCalibrate(unittest.TestCase):
# |
# Resize
G = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 1, 5, 5])
H = helper.make_tensor_value_info("H", TensorProto.FLOAT, [1, 1, 3, 3])
J = helper.make_tensor_value_info("J", TensorProto.FLOAT, [1, 1, 3, 3])
M = helper.make_tensor_value_info("M", TensorProto.FLOAT, [0])
N = helper.make_tensor_value_info("N", TensorProto.FLOAT, [0])
O = helper.make_tensor_value_info("O", TensorProto.FLOAT, [1, 1, 5, 5])
vi_g = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 1, 5, 5])
vi_m = helper.make_tensor_value_info("M", TensorProto.FLOAT, [0])
vi_n = helper.make_tensor_value_info("N", TensorProto.FLOAT, [0])
vi_o = helper.make_tensor_value_info("O", TensorProto.FLOAT, [1, 1, 5, 5])
# O = helper.make_tensor_value_info('O', TensorProto.FLOAT, None)
conv_node_1 = onnx.helper.make_node(
"Conv",
["G", "H"],
["I"],
name="Conv1",
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
"Conv", ["G", "conv1_w"], ["I"], name="Conv1", kernel_shape=[3, 3], pads=[1, 1, 1, 1],
)
conv_node_2 = onnx.helper.make_node(
"Conv",
["I", "J"],
["K"],
name="Conv2",
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
"Conv", ["I", "conv2_w"], ["K"], name="Conv2", kernel_shape=[3, 3], pads=[1, 1, 1, 1],
)
resize_node_1 = onnx.helper.make_node("Resize", ["K", "M", "N"], ["O"], name="Reize1")
graph = helper.make_graph(
[conv_node_1, conv_node_2, resize_node_1],
"test_graph_5",
[G, H, J, M, N],
[O],
)
graph = helper.make_graph([conv_node_1, conv_node_2, resize_node_1], "test_graph_5", [vi_g, vi_m, vi_n], [vi_o],)
conv1_w = generate_input_initializer([1, 1, 3, 3], np.float32, "conv1_w")
conv2_w = generate_input_initializer([1, 1, 3, 3], np.float32, "conv2_w")
graph.initializer.extend([conv1_w, conv2_w])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
test_model_path = "./test_model_5.onnx"
onnx.save(model, test_model_path)
@ -359,25 +362,49 @@ class TestCalibrate(unittest.TestCase):
augmented_model_node_names = [node.name for node in augmented_model.graph.node]
augmented_model_outputs = [output.name for output in augmented_model.graph.output]
added_node_names = [
"G_ReduceMin",
"G_ReduceMax",
"I_ReduceMin",
"I_ReduceMax",
"K_ReduceMin",
"K_ReduceMax",
"M_ReduceMin",
"M_ReduceMax",
"N_ReduceMin",
"N_ReduceMax",
"O_ReduceMin",
"O_ReduceMax",
"G_ReduceMin_Reshape",
"G_ReduceMax_Reshape",
"I_ReduceMin_Reshape",
"I_ReduceMax_Reshape",
"K_ReduceMin_Reshape",
"K_ReduceMax_Reshape",
"M_ReduceMin_Reshape",
"M_ReduceMax_Reshape",
"N_ReduceMin_Reshape",
"N_ReduceMax_Reshape",
"O_ReduceMin_Reshape",
"O_ReduceMax_Reshape",
]
added_outputs = [
"G_ReduceMin",
"G_ReduceMax",
"I_ReduceMin",
"I_ReduceMax",
"K_ReduceMin",
"K_ReduceMax",
"M_ReduceMin",
"M_ReduceMax",
"N_ReduceMin",
"N_ReduceMax",
"O_ReduceMin",
"O_ReduceMax",
]
# Original 3 nodes + added ReduceMin/Max nodes * 8
self.assertEqual(len(augmented_model_node_names), 19)
# Original 1 graph output + added outputs * 8
self.assertEqual(len(augmented_model_outputs), 17)
# Original 3 nodes + (ReduceMin + Reshape, ReduceMax + Reshape) * 6 tensors
self.assertEqual(len(augmented_model_node_names), 27)
# Original 1 graph output + 6 tensors * 2
self.assertEqual(len(augmented_model_outputs), 13)
for name in added_node_names:
self.assertTrue(name in augmented_model_node_names)
for output in added_outputs: