def _convert_quantization_constants_to_initializers(model: ModelProto):
for quant_node in model.graph.node:
if quant_node.op_type not in _QUANTIZE_OP_NAMES:
continue
scale_const = get_nodes_by_output_id(model, quant_node.input[1])
zp_const = get_nodes_by_output_id(model, quant_node.input[2])
consts = scale_const + zp_const
for const in consts:
if const.op_type != "Constant":
continue
# constants should be tensor type
_convert_tensor_constant_to_initializer(model, const)
def get_quantization_params(
model: ModelProto,
node: NodeProto,
include_target: bool = False) -> QuantizationParams:
"""
:param model: ONNX model to read from
:param node: A QuantizeLinear or DequantizeLinear Node
:param include_target: Set True include quantization target. If False,
target value will be returned as None. Default is None
:return: QuantizationParams object with scale and zero point, will include the
quantization target if it is an initializer otherwise target will be None
"""
assert (
node.op_type in _QUANTIZE_OP_NAMES
), "Op Type must be either QuantizeLinear or DequantizeLinear, found {} ".format(
node.op_type)
scale = get_init_by_name(model, node.input[1])
if scale is None:
scale_const = get_nodes_by_output_id(model, node.input[1])
if scale_const:
scale = scale_const[0].attribute[0].t
assert scale, "Quantization scale {} not found".format(node.input[1])
zero_point = get_init_by_name(model, node.input[2])
if zero_point is None:
zero_point_const = get_nodes_by_output_id(model, node.input[2])
if zero_point_const:
zero_point = zero_point_const[0].attribute[0].t
assert zero_point, "Quantization zero point {} not found".format(
node.input[2])
scale = numpy_helper.to_array(scale)
zero_point = numpy_helper.to_array(zero_point)
target = None
if include_target:
target = get_init_by_name(model, node.input[0])
if target is not None:
target = numpy_helper.to_array(target)
return QuantizationParams(scale=scale,
zero_point=zero_point,
target=target)
def test_get_node_by_output_id(simple_onnx_model):
first_node = simple_onnx_model.graph.node[0]
assert get_nodes_by_output_id(simple_onnx_model, "Y") == [first_node]
assert get_nodes_by_output_id(simple_onnx_model, "NONE") == []
def _get_single_node_parent(
model: ModelProto, node: NodeProto, input_idx: int
) -> Union[NodeProto, None]:
# return parent of input node if it only has one parent, otherwise return None
parent = get_nodes_by_output_id(model, node.input[input_idx])
return parent[0] if len(parent) == 1 else None
