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