Skip to main content
`Ctrl`+`K`

# Quantizing ONNX Models with Custom Operators Using Quark

NOTE This tutorial can be downloaded for local execution on a Jupyter Notebook environment. Click here to download the source file.

This tutorial demonstrates how to use Quark to quantize an ONNX model containing Custom Operators. The example includes two types of custom operator implementations: - Python Custom Operator (my_custom_op) - C++ Custom Operator (_COP_IN_OP_NAME)

This tutorial will guide you through the following steps: - Building a floating-point ONNX model with custom operators - Preparing calibration data for quantization - Using Quark to quantize the model - Running inference on the quantized model using ONNX Runtime and printing the output

## Preparing the Floating-Point Model

Before quantization, we first build a floating-point model containing two custom operators: - my_custom_op: A custom operator implemented in Python. - _COP_IN_OP_NAME: A custom operator implemented in C++.

The following code demonstrates how to incorporate these operators when constructing the ONNX graph and generate a floating-point model ready for quantization.

    import copy
    from pathlib import Path

    import numpy as np
    import onnx
    import onnxruntime
    from onnx import helper
    from onnx.onnx_ml_pb2 import TensorProto
    from onnxruntime.quantization import CalibrationDataReader
    from onnxruntime_extensions import PyCustomOpDef, onnx_op
    from onnxruntime_extensions import get_library_path as ext_lib_path

    from quark.onnx import Config, ModelQuantizer
    from quark.onnx.operators.custom_ops import _COP_DOMAIN, _COP_IN_OP_NAME, get_library_path
    from quark.onnx.quantization.config.custom_config import S16S16_MIXED_S8S8_CONFIG

    op_type = "MyCustomOp"
    op_domain = "ai.onnx.contrib"

    @onnx_op(op_type=op_type, domain=op_domain, inputs=[PyCustomOpDef.dt_float], outputs=[PyCustomOpDef.dt_float])
    def my_custom_op(x: np.ndarray[np.dtype[np.float32]]):
        return x * 2

    def prepare_model(float_model_path):
        np.random.seed(123)
        data_type = TensorProto.FLOAT
        data_shape = (1, 3, 5, 5)

        gamma = np.ones(data_shape).astype(np.float32)
        beta = np.zeros(data_shape).astype(np.float32)

        in_param_nodes = [
            helper.make_node(
                "Constant", [], ["gamma"], value=onnx.helper.make_tensor("y_scale", data_type, data_shape, gamma)
            ),
            helper.make_node(
                "Constant", [], ["beta"], value=onnx.helper.make_tensor("y_zero_point", data_type, data_shape, beta)
            ),
        ]

        graph_def = helper.make_graph(
            nodes=[
                helper.make_node(op_type, ["input"], ["input_out"], domain=op_domain),
                helper.make_node(
                    _COP_IN_OP_NAME,
                    ["input_out", "gamma", "beta"],
                    ["y"],
                    domain=_COP_DOMAIN,
                ),
            ]
            + in_param_nodes,
            name="test-in",
            inputs=[helper.make_tensor_value_info("input", data_type, shape=None)],
            outputs=[helper.make_tensor_value_info("y", data_type, shape=None)],
        )

        produce_opset_version = 19
        opset_imports = [onnx.helper.make_operatorsetid("", produce_opset_version)]
        model_def = helper.make_model(graph_def, producer_name="quark.onnx", ir_version=9, opset_imports=opset_imports)
        onnx.save(model_def, float_model_path)
        print(f"Model has been saved to {float_model_path}")

    float_model_path = Path("user_custom_op_model.onnx").as_posix()
    prepare_model(float_model_path)

## Preparing Calibration Data

To demonstrate the quantization process, a set of pseudo input data with the shape `(1, 3, 5, 5)` is constructed as calibration samples. The calibration data is provided through a custom DataReader class, which serves as the standard input data interface for the Quark quantizer.

    input_tensor = np.array(
        [
            [
                [
                    [0.39250988, 0.34032542, 0.91402656, 0.40040675, 0.39050988],
                    [0.39150988, 0.07389439, 0.38661167, 0.8645387, 0.55553377],
                    [0.21639073, 0.10061733, 0.19072777, 0.32449463, 0.79694337],
                    [0.66819394, 0.03191912, 0.397995, 0.01690937, 0.63425934],
                    [0.37730116, 0.80095553, 0.77266306, 0.54853624, 0.27609143],
                ],
                [
                    [0.8222164, 0.5256697, 0.2953402, 0.47371042, 0.40800324],
                    [0.6019997, 0.7506883, 0.5605579, 0.7274801, 0.19008774],
                    [0.76555413, 0.6223917, 0.27387974, 0.85017425, 0.70976704],
                    [0.868642, 0.18798842, 0.26945123, 0.8975411, 0.1434885],
                    [0.30794197, 0.13901855, 0.8121448, 0.8238567, 0.33238393],
                ],
                [
                    [0.57792556, 0.98300576, 0.8607786, 0.6592352, 0.22613065],
                    [0.7223881, 0.46592003, 0.3890724, 0.868129, 0.691695],
                    [0.4210194, 0.5127264, 0.6360194, 0.30745587, 0.1583932],
                    [0.67081225, 0.16967775, 0.6681447, 0.71011454, 0.3408417],
                    [0.83913565, 0.3341194, 0.8299601, 0.9870858, 0.35757536],
                ],
            ]
        ],
        dtype=np.float32,
    )

    class DataReader(CalibrationDataReader):
        def __init__(self, input_tensor):
            self.data = [input_tensor]
            self.input_name = "input"
            self.index = 0

        def get_next(self):
            if self.index < len(self.data):
                input_dict = {self.input_name: self.data[self.index]}
                self.index += 1
                return input_dict
            else:
                return None

        def rewind(self):
            self.index = 0

    data_reader = DataReader(input_tensor)

## Quantizing the Model

In this step, we use the prepared floating-point model and calibration data, and perform quantization by configuring and executing it with the `quantize_model` function. Quark will automatically handle custom operators.

    def quantize_model(float_model_path, quantized_model_path, data_reader):
        config_copy = copy.deepcopy(S16S16_MIXED_S8S8_CONFIG)
        config_copy.extra_options["UserCustomOpLibPath"] = [get_library_path(), ext_lib_path()]
        quant_config = Config(global_quant_config=config_copy)
        quantizer = ModelQuantizer(quant_config)
        quantizer.quantize_model(float_model_path, quantized_model_path, data_reader)
        print("Quantized the ONNX model and saved it at:", quantized_model_path)

    quantized_model_path = Path("user_custom_op_model_quantized.onnx").as_posix()
    quantize_model(float_model_path, quantized_model_path, data_reader)

## Inference with the Quantized Model

After quantization, we use ONNX Runtime to load the quantized model and perform inference. The model output is printed to verify that the quantized model runs successfully and that custom operators are correctly applied.

    def infer_quantized_model(quantized_model_path):
        sess_options = onnxruntime.SessionOptions()
        sess_options.register_custom_ops_library(ext_lib_path())
        sess_options.register_custom_ops_library(get_library_path())
        sess = onnxruntime.InferenceSession(quantized_model_path, sess_options)
        input_name = sess.get_inputs()[0].name
        output_name = sess.get_outputs()[0].name
        input_data = input_tensor
        output = sess.run([output_name], {input_name: input_data})
        print(f"Model output: {output}")

    infer_quantized_model(quantized_model_path)

## Summary

This guide demonstrated how to: - Build an ONNX model with custom operators (Python & C++) - Prepare calibration data for quantization - Quantize the model using Quark - Perform inference on the quantized model using ONNX Runtime

With these steps, you can seamlessly apply Quark to any ONNX model containing custom operators, making quantization more flexible and controllable.