Support per-channel quantization of weight tensor (#5057)

* Support per-channel quantization of weight tensor

* rename util functions

* fix bugs in calibrate

* add support of reduce_range

* refine opset check
This commit is contained in:
Yufeng Li 2020-09-14 11:53:50 -07:00 committed by GitHub
parent 2a456d16c0
commit 20b2f45b24
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 148 additions and 126 deletions

View file

@ -56,6 +56,7 @@ 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 = []
@ -167,7 +168,7 @@ class ONNXCalibrater:
else:
self.input_name_to_nodes[input_name].append(node)
def calculate_scale_zeropoint(self, node, next_node, rmin, rmax):
def calculate_scale_zeropoint(self, next_node, rmin, rmax):
zp_and_scale = []
# adjust rmin and rmax such that 0 is included in the range. This is required
@ -178,16 +179,17 @@ class ONNXCalibrater:
# We update the output range min and max when next node is clip or relu
# With this technique we can remove these 2 ops and
# reduce the output range which in turn helps to improve accuracy
if next_node.op_type == 'Clip':
clip_min = next_node.attribute[0].f
clip_max = next_node.attribute[1].f
if rmin < clip_min:
rmin = clip_min
if rmax > clip_max:
rmax = clip_max
if next_node.op_type == 'Relu':
if rmin < 0:
rmin = 0
if next_node:
if next_node.op_type == 'Clip':
clip_min = next_node.attribute[0].f
clip_max = next_node.attribute[1].f
if rmin < clip_min:
rmin = clip_min
if rmax > clip_max:
rmax = clip_max
elif next_node.op_type == 'Relu':
if rmin < 0:
rmin = 0
scale = np.float32((rmax - rmin) / 255 if rmin != rmax else 1)
initial_zero_point = (0 - rmin) / scale
@ -227,16 +229,15 @@ class ONNXCalibrater:
self._get_input_name_to_nodes(model)
for node in model.graph.node:
for node_output_name in node.output:
if node_output_name in self.input_name_to_nodes:
children = self.input_name_to_nodes[node_output_name]
for child in children:
if node_output_name in quantization_thresholds:
node_thresholds = quantization_thresholds[node_output_name]
node_params = self.calculate_scale_zeropoint(node, child, node_thresholds[0],
node_thresholds[1])
quantization_params[node_output_name] = node_params
for tensor_name in quantization_thresholds.keys():
child = None
if tensor_name in self.input_name_to_nodes:
children = self.input_name_to_nodes[tensor_name]
if (len(children) == 1):
child = children[0]
node_thresholds = quantization_thresholds[tensor_name]
node_params = self.calculate_scale_zeropoint(child, node_thresholds[0], node_thresholds[1])
quantization_params[tensor_name] = node_params
return quantization_params

View file

@ -1,5 +1,5 @@
import onnx
from .quant_utils import _find_by_name
from .quant_utils import find_by_name
class ONNXModel:
@ -37,7 +37,7 @@ class ONNXModel:
self.model.graph.node.extend(nodes_to_add)
def add_initializer(self, tensor):
if _find_by_name(tensor.name, self.model.graph.initializer) is None:
if find_by_name(tensor.name, self.model.graph.initializer) is None:
self.model.graph.initializer.extend([tensor])
def get_initializer(self, name):
@ -112,7 +112,7 @@ class ONNXModel:
'''
graph_nodes_list = list(graph.node) #deep copy
graph_nodes_list.extend(new_nodes_list)
node = _find_by_name(node_name, graph_nodes_list)
node = find_by_name(node_name, graph_nodes_list)
return node
def find_nodes_by_initializer(self, graph, initializer):

View file

@ -16,7 +16,7 @@ from onnx import shape_inference
from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel
from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue, quantization_modes
from .quant_utils import _find_by_name, _get_elem_index, _get_mul_node, _generate_identified_filename, _attribute_to_kwarg
from .quant_utils import find_by_name, get_elem_index, get_mul_node, generate_identified_filename, attribute_to_kwarg
from .quant_utils import QuantType, onnx_domain, __producer__, __version__
from .registry import CreateOpQuantizer, CreateDefaultOpQuantizer
@ -46,7 +46,7 @@ def quantize_data(data, quantize_range, qType):
if qType == onnx_proto.TensorProto.INT8:
max_range = max(abs(rmin), abs(rmax))
scale = (float(max_range) * 2) / quantize_range
scale = (float(max_range) * 2) / quantize_range if max_range > 0 else 1
zero_point = 0
# signed byte type
quantized_data = (np.asarray(data) / scale).round().astype('b')
@ -60,27 +60,28 @@ def quantize_data(data, quantize_range, qType):
return rmin, rmax, zero_point, scale, quantized_data
def _get_qrange_for_qType(qType):
def _get_qrange_for_qType(qType, reduce_range=False):
'''
Helper function to get the quantization range for a type.
parameter qType: quantization type.
return: quantization range.
'''
if qType == onnx_proto.TensorProto.UINT8:
return 255 # 2^b - 1
return 127 if reduce_range else 255
elif qType == onnx_proto.TensorProto.INT8:
return 254 # [-(2^{b-1}-1), 2^{b-1}-1]: [-127, 127] for 8 bits.
return 128 if reduce_range else 254 # [-64, 64] for reduce_range, and [-127, 127] full_range.
else:
raise ValueError('unsupported quantization data type')
class ONNXQuantizer:
def __init__(self, model, per_channel, mode, static, weight_qType, input_qType, quantization_params,
def __init__(self, model, per_channel, reduce_range, mode, static, weight_qType, input_qType, quantization_params,
nodes_to_quantize, nodes_to_exclude, op_types_to_quantize):
onnx_model = shape_inference.infer_shapes(model)
self.model = ONNXModel(onnx_model)
self.value_infos = {vi.name: vi for vi in onnx_model.graph.value_info}
self.per_channel = per_channel # weight-pack per channel
self.reduce_range = reduce_range
self.mode = mode # QuantizationMode.Value
self.static = static # use static quantization for inputs.
self.fuse_dynamic_quant = False
@ -119,15 +120,18 @@ class ONNXQuantizer:
raise ValueError('Failed to find proper ai.onnx domain')
opset_version = ai_onnx_domain[0].version
if opset_version < 10:
raise ValueError("The original model opset version is {}, which does not support quantized operators.\n\
The opset version of quantized model will be set to 10. Use onnx model checker to verify model after quantization."
.format(opset_version))
if opset_version == 10:
self.fuse_dynamic_quant = False
else:
self.fuse_dynamic_quant = True
print(
"Warning: The original model opset version is {}, which does not support node fusions. Please update the model to opset >= 11 for better performance."
.format(opset_version))
return
if opset_version < 10:
print(
"Warning: The original model opset version is {}, which does not support quantization. Please update the model to opset >= 11. Updating the model automatically to opset 11. Please verify the quantized model."
.format(opset_version))
self.model.model.opset_import.remove(ai_onnx_domain[0])
self.model.model.opset_import.extend([onnx.helper.make_opsetid("", 11)])
def replace_gemm_with_matmul(self):
nodes_to_remove = []
@ -196,8 +200,8 @@ class ONNXQuantizer:
# TODO: convert it to the specified input_type
scale_tensor_name = curr_node.input[1]
zp_tensor_name = curr_node.input[2]
initializer_scale = _find_by_name(scale_tensor_name, self.model.initializer())
initializer_zp = _find_by_name(zp_tensor_name, self.model.initializer())
initializer_scale = find_by_name(scale_tensor_name, self.model.initializer())
initializer_zp = find_by_name(zp_tensor_name, self.model.initializer())
zp_and_scale = [
onnx.numpy_helper.to_array(initializer_zp),
onnx.numpy_helper.to_array(initializer_scale)
@ -205,7 +209,7 @@ class ONNXQuantizer:
#connect the previous and successive node input and output
for succ_node in succ_nodes:
succ_idx = _get_elem_index(next_node.output[0], succ_node.input)
succ_idx = get_elem_index(next_node.output[0], succ_node.input)
if succ_idx != -1:
succ_node.input[succ_idx] = curr_node.input[0]
else:
@ -273,11 +277,8 @@ class ONNXQuantizer:
return self.model.model
def find_weight_data(self, initializer):
'''
:param initializer: TensorProto initializer object from a graph
:return: a list of initialized data in a given initializer object
'''
@staticmethod
def tensor_proto_to_array(initializer):
if initializer.data_type == onnx_proto.TensorProto.FLOAT:
weights = onnx.numpy_helper.to_array(initializer)
else:
@ -285,6 +286,10 @@ class ONNXQuantizer:
initializer.name, type_to_name[initializer.data_type]))
return weights
def is_input_a_weight(self, input_name):
initializer = find_by_name(input_name, self.model.initializer())
return initializer is not None
def _is_valid_quantize_value(self, value_name):
if value_name in self.value_infos:
value_info = self.value_infos[value_name]
@ -293,11 +298,11 @@ class ONNXQuantizer:
return self._is_valid_initializer_value(value_name)
def _is_valid_initializer_value(self, value_name):
weight = _find_by_name(value_name, self.model.initializer())
weight = find_by_name(value_name, self.model.initializer())
return weight is not None and weight.data_type == onnx_proto.TensorProto.FLOAT
def _is_valid_quantize_weight(self, weight_name):
weight = _find_by_name(weight_name, self.model.initializer())
weight = find_by_name(weight_name, self.model.initializer())
return weight is not None and weight.data_type == onnx_proto.TensorProto.FLOAT
def _remove_quantized_weights(self):
@ -318,7 +323,7 @@ class ONNXQuantizer:
if self.model.ir_version() < 4:
print("Warning: invalid weight name {} found in the graph (not a graph input)".format(weight.name))
def _update_graph(self, weight):
def _update_weight(self, weight):
'''
Given a weight object, update the graph by doing the following:
- remove old initializer, update new initializers for quantized weight, zero point, and scale
@ -357,9 +362,9 @@ class ONNXQuantizer:
:param qType: type to quantize to
:return: Weight class with quantization information
'''
weights_data = self.find_weight_data(initializer)
rmin, rmax, zero_point, scale, quantized_weights_data = quantize_data(weights_data.flatten().tolist(),
_get_qrange_for_qType(qType), qType)
weights_data = self.tensor_proto_to_array(initializer)
rmin, rmax, zero_point, scale, quantized_weights_data = quantize_data(
weights_data.flatten().tolist(), _get_qrange_for_qType(qType, self.reduce_range), qType)
weight = QuantizedInitializer(initializer.name,
initializer, [rmin], [rmax], [zero_point], [scale],
weights_data,
@ -375,7 +380,7 @@ class ONNXQuantizer:
return weight
def _get_quantized_weight_convolution(self, initializer, qType):
def _get_quantized_weight_per_channel(self, initializer, qType, channel_axis):
'''
:param initializer: initializer TypeProto to quantize
:param qType: type to quantize to
@ -384,38 +389,34 @@ class ONNXQuantizer:
if not self.per_channel:
return self._get_quantized_weight(initializer, qType)
weights = self.find_weight_data(initializer)
# Quantize per output channel
# Assuming (M x C/group x kH x kW) format where M is number of output channels.
channel_count = initializer.dims[0]
np_data = np.reshape(weights, initializer.dims)
weights = self.tensor_proto_to_array(initializer)
channel_count = weights.shape[channel_axis]
rmin_list = []
rmax_list = []
zero_point_list = []
scale_list = []
quantized_per_channel_data_list = []
for i in range(channel_count):
# for each channel, compute quantization data. Assuming (M x C/group x kH x kW)
per_channel_data = np_data[i, :, :, :].flatten()
per_channel_data = weights.take(i, channel_axis)
rmin, rmax, zero_point, scale, quantized_per_channel_data = quantize_data(
per_channel_data.flatten().tolist(), _get_qrange_for_qType(qType), qType)
per_channel_data.flatten().tolist(), _get_qrange_for_qType(qType, self.reduce_range), qType)
rmin_list.append(rmin)
rmax_list.append(rmax)
zero_point_list.append(zero_point)
scale_list.append(scale)
quantized_per_channel_data_list.append(quantized_per_channel_data)
channel_index = 0 # (M x C/group x kH x kW)
# combine per_channel_data into one
reshape_dims = list(initializer.dims) # deep copy
reshape_dims[channel_index] = 1 # only one per channel for reshape
reshape_dims = list(weights.shape) # deep copy
reshape_dims[channel_axis] = 1 # only one per channel for reshape
quantized_weights = np.asarray(quantized_per_channel_data_list[0]).reshape(reshape_dims)
for i in range(1, len(quantized_per_channel_data_list)):
channel_weights = np.asarray(quantized_per_channel_data_list[i]).reshape(reshape_dims)
quantized_weights = np.concatenate((quantized_weights, channel_weights), axis=0)
quantized_weights = np.concatenate((quantized_weights, channel_weights), channel_axis)
weight = QuantizedInitializer(initializer.name, initializer, rmin_list, rmax_list, zero_point_list, scale_list,
weights,
quantized_weights.flatten().tolist(), channel_index, qType)
quantized_weights.flatten().tolist(), channel_axis, qType)
# Make entry for this quantized weight
assert (weight.name not in self.quantized_value_map)
@ -642,7 +643,7 @@ class ONNXQuantizer:
return nodes + [qlinear_node]
def _get_bias_add_nodes(self, nodes, node, last_output, quantized_bias_name):
def get_bias_add_nodes(self, nodes, node, last_output, quantized_bias_name):
'''
Given a node, this function handles bias add by adding a "reshape" node on bias and an "add" node
parameter nodes: new nodes would be appended into nodes
@ -738,13 +739,13 @@ class ONNXQuantizer:
# get scale for weight
weight_scale_name = self.quantized_value_map[node.input[1]].scale_name
weight_initializer = _find_by_name(weight_scale_name, self.model.initializer())
weight_scale = self.find_weight_data(weight_initializer)
weight_initializer = find_by_name(weight_scale_name, self.model.initializer())
weight_scale = self.tensor_proto_to_array(weight_initializer)
# get bias
bias_name = node.input[2]
bias_initializer = _find_by_name(bias_name, self.model.initializer())
bias_data = self.find_weight_data(bias_initializer)
bias_initializer = find_by_name(bias_name, self.model.initializer())
bias_data = self.tensor_proto_to_array(bias_initializer)
quantized_bias_name = bias_name + "_quantized"
# input scale is not provided and this input is dynamically quantized so it is not pre-computed at this point
@ -762,8 +763,8 @@ class ONNXQuantizer:
raise ValueError("Expected {} to be in quantized value map for static quantization".format(
node.input[0]))
inputscale_initializer = _find_by_name(input_scale_name, self.model.initializer())
input_scale = self.find_weight_data(inputscale_initializer)
inputscale_initializer = find_by_name(input_scale_name, self.model.initializer())
input_scale = self.tensor_proto_to_array(inputscale_initializer)
# calcuate scale for bias
@ -792,7 +793,7 @@ class ONNXQuantizer:
return quantized_bias_name
def _quantize_inputs(self, node, indices):
def quantize_inputs(self, node, indices):
'''
Given a node, this function quantizes the inputs as follows:
- If input is an initializer, quantize the initializer data, replace old initializer
@ -800,17 +801,15 @@ class ONNXQuantizer:
- Else, add QuantizeLinear nodes to perform quantization
parameter node: node being quantized in NodeProto format.
parameter indices: input indices to quantize.
parameter new_nodes_list: List of new nodes created before processing this node. This is used to
check that two QuantizeLinear nodes are not being added for same input.
return: (List of quantized input names,
List of zero point names used for input quantization,
List of scale names used for input quantization,
List of new QuantizeLinear nodes created)
'''
quantized_input_names = []
zero_point_names = []
scale_names = []
zero_point_names = []
quantized_input_names = []
nodes = []
for input_index in indices:
@ -819,27 +818,18 @@ class ONNXQuantizer:
# Find if this input is already quantized
if node_input in self.quantized_value_map:
quantized_value = self.quantized_value_map[node_input]
qType = self.weight_qType if quantized_value.value_type == QuantizedValueType.Initializer else self.input_qType
if quantized_value.qType != qType:
raise ValueError(
"{} is being used by multiple nodes which are being quantized to different types. "
"This is not suported.", node_input)
quantized_input_names.append(quantized_value.q_name)
scale_names.append(quantized_value.scale_name)
zero_point_names.append(quantized_value.zp_name)
quantized_input_names.append(quantized_value.q_name)
continue
# Quantize the input
initializer = _find_by_name(node_input, self.model.initializer())
initializer = find_by_name(node_input, self.model.initializer())
if initializer is not None:
if node.op_type == "Conv":
weight = self._get_quantized_weight_convolution(initializer, self.weight_qType)
else:
weight = self._get_quantized_weight(initializer, self.weight_qType)
weight = self._get_quantized_weight(initializer, self.weight_qType)
# Update graph
self._update_graph(weight)
self._update_weight(weight)
quantized_input_names.append(weight.name + "_quantized")
zero_point_names.append(weight.name + "_zero_point")
@ -864,6 +854,20 @@ class ONNXQuantizer:
return (quantized_input_names, zero_point_names, scale_names, nodes)
def quantize_weight_per_channel(self, weight_name, axis):
# Find if this input is already quantized
if weight_name in self.quantized_value_map:
quantized_value = self.quantized_value_map[weight_name]
return (quantized_value.zp_name, quantized_value.q_name, quantized_value.scale_name)
else:
initializer = find_by_name(weight_name, self.model.initializer())
if initializer is None:
raise ValueError("{} is not an initializer", weight_name)
weight = self._get_quantized_weight_per_channel(initializer, self.weight_qType, axis)
self._update_weight(weight)
return (weight.name + "_quantized", weight.name + "_zero_point", weight.name + "_scale")
def _dequantize_value(self, value_name):
'''
Given a value (input/output) which is quantized, add a DequantizeLinear node to dequantize

View file

@ -1,6 +1,6 @@
import onnx
from .base_operator import QuantOperatorBase
from ..quant_utils import _attribute_to_kwarg, ms_domain
from ..quant_utils import attribute_to_kwarg, ms_domain
from onnx import onnx_pb as onnx_proto
'''
Quantize Attention
@ -21,7 +21,7 @@ class AttentionQuant(QuantOperatorBase):
assert (node.op_type == "Attention")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer._quantize_inputs(node, [0, 1])
self.quantizer.quantize_inputs(node, [0, 1])
qattention_name = "" if node.name == "" else node.name + "_quant"
@ -35,7 +35,7 @@ class AttentionQuant(QuantOperatorBase):
kwargs = {}
for attribute in node.attribute:
kwargs.update(_attribute_to_kwarg(attribute))
kwargs.update(attribute_to_kwarg(attribute))
kwargs["domain"] = ms_domain
qattention_node = onnx.helper.make_node("QAttention", inputs, node.output, qattention_name, **kwargs)
nodes.append(qattention_node)

View file

@ -1,6 +1,6 @@
import onnx
from .base_operator import QuantOperatorBase
from ..quant_utils import _attribute_to_kwarg, ms_domain, QuantizedValue, QuantizedValueType
from ..quant_utils import attribute_to_kwarg, ms_domain, QuantizedValue, QuantizedValueType
from onnx import onnx_pb as onnx_proto
@ -17,14 +17,14 @@ class QLinearBinaryOp(QuantOperatorBase):
return super().quantize()
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer._quantize_inputs(node, [0, 1])
self.quantizer.quantize_inputs(node, [0, 1])
qlinear_binary_math_output = node.output[0] + "_quantized"
qlinear_binary_math_name = node.name + "_quant" if node.name != "" else ""
kwargs = {}
for attribute in node.attribute:
kwargs.update(_attribute_to_kwarg(attribute))
kwargs.update(attribute_to_kwarg(attribute))
kwargs["domain"] = ms_domain
qlinear_binary_math_inputs = []

View file

@ -1,6 +1,6 @@
import onnx
from .base_operator import QuantOperatorBase
from ..quant_utils import _find_by_name, _get_mul_node, QuantizedValue, QuantizedValueType, _attribute_to_kwarg
from ..quant_utils import find_by_name, get_mul_node, QuantizedValue, QuantizedValueType, attribute_to_kwarg
from onnx import onnx_pb as onnx_proto
@ -13,7 +13,7 @@ class ConInteger(QuantOperatorBase):
assert (node.op_type == "Conv")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer._quantize_inputs(node, [0, 1])
self.quantizer.quantize_inputs(node, [0, 1])
# quantize bias if exist
quantized_bias_name = ""
@ -27,7 +27,7 @@ class ConInteger(QuantOperatorBase):
kwargs = {}
for attribute in node.attribute:
kwargs.update(_attribute_to_kwarg(attribute))
kwargs.update(attribute_to_kwarg(attribute))
conv_integer_node = onnx.helper.make_node("ConvInteger", quantized_input_names + zero_point_names,
[conv_integer_output], conv_integer_name, **kwargs)
nodes.append(conv_integer_node)
@ -51,9 +51,9 @@ class ConInteger(QuantOperatorBase):
else:
scales_mul_op = scale_names[0] + "_" + scale_names[1] + "_mul"
scales_mul_node = _find_by_name(scales_mul_op, self.nodes)
scales_mul_node = find_by_name(scales_mul_op, self.quantizer.new_nodes)
if scales_mul_node is None:
scales_mul_node = _get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op)
scales_mul_node = get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op)
nodes.append(scales_mul_node)
scales_mul_op_output = scales_mul_node.output[0]
@ -61,9 +61,9 @@ class ConInteger(QuantOperatorBase):
# Add mul operation to multiply mul_scales_op result with output of ConvInteger
# and make the output of this node the same as output of original conv node.
output_scale_mul_op = conv_integer_name + "_output_scale_mul" if conv_integer_name != "" else ""
nodes.append(_get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op))
nodes.append(get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op))
self.new_nodes += nodes
self.quantizer.new_nodes += nodes
class QLinearCov(QuantOperatorBase):
@ -74,8 +74,16 @@ class QLinearCov(QuantOperatorBase):
node = self.node
assert (node.op_type == "Conv")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer._quantize_inputs(node, [0, 1])
if self.quantizer.is_input_a_weight(node.input[1]):
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer.quantize_inputs(node, [0])
quant_weight_tuple = self.quantizer.quantize_weight_per_channel(node.input[1], 0)
quantized_input_names.append(quant_weight_tuple[0])
zero_point_names.append(quant_weight_tuple[1])
scale_names.append(quant_weight_tuple[2])
else:
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer.quantize_inputs(node, [0, 1])
quantized_bias_name = ""
bias_present = False
@ -94,7 +102,7 @@ class QLinearCov(QuantOperatorBase):
kwargs = {}
for attribute in node.attribute:
kwargs.update(_attribute_to_kwarg(attribute))
kwargs.update(attribute_to_kwarg(attribute))
qlinear_conv_inputs = []
# Input 0
qlinear_conv_inputs.append(quantized_input_names[0])

View file

@ -15,7 +15,7 @@ class EmbedLayerNormalizationQuant(QuantOperatorBase):
assert (node.op_type == "EmbedLayerNormalization")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer._quantize_inputs(node, [2, 3, 4])
self.quantizer.quantize_inputs(node, [2, 3, 4])
nodes.append(node)

View file

@ -19,7 +19,7 @@ class GatherQuant(QuantOperatorBase):
return
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer._quantize_inputs(node, [0])
self.quantizer.quantize_inputs(node, [0])
gather_new_output = node.output[0] + "_quantized"

View file

@ -1,6 +1,6 @@
import onnx
from .base_operator import QuantOperatorBase
from ..quant_utils import _find_by_name, _get_mul_node, QuantizedValue, QuantizedValueType
from ..quant_utils import find_by_name, get_mul_node, QuantizedValue, QuantizedValueType
from onnx import onnx_pb as onnx_proto
'''
Used when quantize mode is QuantizationMode.IntegerOps.
@ -16,7 +16,7 @@ class MatMulInteger(QuantOperatorBase):
assert (node.op_type == "MatMul")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer._quantize_inputs(node, [0, 1])
self.quantizer.quantize_inputs(node, [0, 1])
matmul_integer_output = node.output[0] + "_quantized"
matmul_integer_name = node.name + "_quant" if node.name != "" else ""
@ -36,9 +36,9 @@ class MatMulInteger(QuantOperatorBase):
scales_mul_op = matmul_integer_name + "_scales_mul" if matmul_integer_name != "" else scale_names[
0] + "_" + scale_names[1] + "_mul"
scales_mul_node = _find_by_name(scales_mul_op, self.quantizer.new_nodes)
scales_mul_node = find_by_name(scales_mul_op, self.quantizer.new_nodes)
if scales_mul_node is None:
scales_mul_node = _get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op)
scales_mul_node = get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op)
nodes.append(scales_mul_node)
scales_mul_op_output = scales_mul_node.output[0]
@ -48,7 +48,7 @@ class MatMulInteger(QuantOperatorBase):
output_scale_mul_op = ""
if matmul_integer_name != "":
output_scale_mul_op = matmul_integer_name + "_output_scale_mul"
nodes.append(_get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op))
nodes.append(get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op))
self.quantizer.new_nodes += nodes
@ -66,7 +66,7 @@ class QLinearMatMul(QuantOperatorBase):
assert (node.op_type == "MatMul")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer._quantize_inputs(node, [0, 1])
self.quantizer.quantize_inputs(node, [0, 1])
data_found, output_scale_name, output_zp_name, _, _ = \
self.quantizer._get_quantization_params(node.output[0])

View file

@ -104,7 +104,7 @@ class QuantizedValue:
self.qType = qType
def _attribute_to_kwarg(attribute):
def attribute_to_kwarg(attribute):
'''
Convert attribute to kwarg format for use with onnx.helper.make_node.
:parameter attribute: attribute in AttributeProto format.
@ -141,7 +141,7 @@ def _attribute_to_kwarg(attribute):
return {attribute.name: value}
def _find_by_name(item_name, item_list):
def find_by_name(item_name, item_list):
'''
Helper function to find item by name in a list.
parameter item_name: name of the item.
@ -152,7 +152,7 @@ def _find_by_name(item_name, item_list):
return items[0] if len(items) > 0 else None
def _get_elem_index(elem_name, elem_list):
def get_elem_index(elem_name, elem_list):
'''
Helper function to return index of an item in a node list
'''
@ -163,7 +163,7 @@ def _get_elem_index(elem_name, elem_list):
return elem_idx
def _get_mul_node(inputs, output, name):
def get_mul_node(inputs, output, name):
'''
Helper function to create a Mul node.
parameter inputs: list of input names.
@ -174,7 +174,7 @@ def _get_mul_node(inputs, output, name):
return onnx.helper.make_node("Mul", inputs, [output], name)
def _generate_identified_filename(filename: Path, identifier: str) -> Path:
def generate_identified_filename(filename: Path, identifier: str) -> Path:
'''
Helper function to generate a identifiable filepath by concatenating the given identifier as a suffix.
'''

View file

@ -16,7 +16,7 @@ from onnx import shape_inference
from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel
from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue, quantization_modes
from .quant_utils import _find_by_name, _get_elem_index, _get_mul_node, _generate_identified_filename, _attribute_to_kwarg
from .quant_utils import find_by_name, get_elem_index, get_mul_node, generate_identified_filename, attribute_to_kwarg
from .quant_utils import QuantType
from .registry import CreateOpQuantizer, CreateDefaultOpQuantizer, QLinearOpsRegistry, IntegerOpsRegistry
@ -32,7 +32,7 @@ def optimize_model(model_path: Path):
parameter model_path: path to the original onnx model
return: optimized onnx model
'''
opt_model_path = _generate_identified_filename(model_path, "-opt")
opt_model_path = generate_identified_filename(model_path, "-opt")
sess_option = SessionOptions()
sess_option.optimized_model_filepath = opt_model_path.as_posix()
sess_option.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC
@ -126,6 +126,7 @@ def quantize_static(model_input,
calibration_data_reader: CalibrationDataReader,
op_types_to_quantize=[],
per_channel=False,
reduce_range=False,
activation_type=QuantType.QUInt8,
weight_type=QuantType.QUInt8,
nodes_to_quantize=[],
@ -138,6 +139,7 @@ def quantize_static(model_input,
:param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default.
:param op_types: operators to quantize
:param per_channel: quantize weights per channel
:param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode
:param activation_type: quantization data type of activation
:param weight_type: quantization data type of weight
:param nodes_to_quantize:
@ -153,8 +155,8 @@ def quantize_static(model_input,
when it is not None.
'''
if activation_type != QuantType.QUInt8 or weight_type != QuantType.QUInt8:
raise ValueError("Static quantization only support uint8 now.")
if activation_type != QuantType.QUInt8:
raise ValueError("Static quantization only support uint8 for activation now.")
input_qType = onnx_proto.TensorProto.INT8 if activation_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
weight_qType = onnx_proto.TensorProto.INT8 if weight_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
@ -169,6 +171,7 @@ def quantize_static(model_input,
quantizer = ONNXQuantizer(
onnx.load(model_input),
per_channel,
reduce_range,
mode,
True, # static
weight_qType,
@ -186,6 +189,7 @@ def quantize_dynamic(model_input: Path,
model_output: Path,
op_types_to_quantize=[],
per_channel=False,
reduce_range=False,
activation_type=QuantType.QUInt8,
weight_type=QuantType.QUInt8,
nodes_to_quantize=[],
@ -194,8 +198,9 @@ def quantize_dynamic(model_input: Path,
Given an onnx model, create a quantized onnx model and save it into a file
:param model_input: file path of model to quantize
:param model_output: file path of quantized model
:param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default.
:param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default
:param per_channel: quantize weights per channel
:param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode
:param nbits: number of bits to represent quantized data. Currently only supporting 8-bit types
:param activation_type: quantization data type of activation
:param weight_type: quantization data type of weight
@ -225,6 +230,7 @@ def quantize_dynamic(model_input: Path,
quantizer = ONNXQuantizer(
optimized_model,
per_channel,
reduce_range,
mode,
False, #static
weight_qType,
@ -242,6 +248,7 @@ def quantize_qat(model_input: Path,
model_output: Path,
op_types_to_quantize=[],
per_channel=False,
reduce_range=False,
activation_type=QuantType.QUInt8,
weight_type=QuantType.QUInt8,
nodes_to_quantize=[],
@ -250,8 +257,9 @@ def quantize_qat(model_input: Path,
Given a quantize-aware traning onnx model, create a quantized onnx model and save it into a file
:param model_input: file path of model to quantize
:param model_output: file path of quantized model
:param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default.
:param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default
:param per_channel: quantize weights per channel
:param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode
:param activation_type: quantization data type of activation
:param nodes_to_quantize:
List of nodes names to quantize. When this list is not None only the nodes in this list
@ -279,6 +287,7 @@ def quantize_qat(model_input: Path,
quantizer = ONNXQuantizer(
optimized_model,
per_channel,
reduce_range,
mode,
False, #static
weight_qType,