tools/quantize.py

"""Rebuild the published dynamic 8-bit Uni2005 ONNX variants."""

from pathlib import Path

import onnx
from onnxconverter_common import float16
from onnxruntime.quantization import QuantType, quantize_dynamic


ROOT = Path(__file__).resolve().parents[1]
SOURCE = ROOT / "onnx" / "model_fp32.onnx"


def quantize(output_name: str, weight_type: QuantType) -> None:
    quantize_dynamic(
        model_input=SOURCE,
        model_output=ROOT / "onnx" / output_name,
        op_types_to_quantize=["LSTM", "MatMul"],
        weight_type=weight_type,
    )


if __name__ == "__main__":
    fp32 = onnx.load(SOURCE)
    fp16 = float16.convert_float_to_float16(
        fp32,
        keep_io_types=True,
        node_block_list=["/Cast"],
    )
    onnx.save(fp16, ROOT / "onnx" / "model_fp16.onnx")
    quantize("model_quantized.onnx", QuantType.QInt8)
    quantize("model_uint8.onnx", QuantType.QUInt8)