File size: 898 Bytes
52fa8e3 004577e 52fa8e3 004577e 52fa8e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """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)
|