quantize_cached.py

import os
from onnxruntime.quantization import quantize_dynamic, QuantType
OUT='/work/out'
for name in ('decoder_init', 'decoder_step'):
    src, dst = f'{OUT}/{name}.onnx', f'{OUT}/{name}.int8.onnx'
    quantize_dynamic(src, dst, weight_type=QuantType.QInt8, extra_options={'MatMulConstBOnly': True})
    fp32 = os.path.getsize(src) + (os.path.getsize(src + '.data') if os.path.exists(src + '.data') else 0)
    print(f'{name}: fp32 {fp32/1e6:.1f} MB -> int8 {os.path.getsize(dst)/1e6:.1f} MB')