| """int8 (dynamic, MatMulInteger) and int4 (MatMulNBits + GatherBlockQuantized) from an fp32 export.""" |
| import sys, os, time |
| import onnx |
| src, outdir = sys.argv[1], sys.argv[2] |
| which = sys.argv[3:] or ["int8", "int4"] |
| clean = os.path.join(outdir, "clean.onnx") |
| if not os.path.exists(clean): |
| m0 = onnx.load(src) |
| del m0.graph.value_info[:] |
| onnx.save(m0, clean, save_as_external_data=True, location="clean.onnx.data") |
| del m0 |
| src = clean |
|
|
| if "int8" in which: |
| from onnxruntime.quantization import quantize_dynamic, QuantType |
| t = time.time() |
| quantize_dynamic(src, os.path.join(outdir, "model_int8.onnx"), weight_type=QuantType.QInt8, |
| per_channel=True, op_types_to_quantize=["MatMul", "Gemm"], |
| extra_options={"MatMulConstBOnly": True}) |
| print("int8 %.0fs" % (time.time() - t)) |
|
|
| if "int4" in which: |
| from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer, DefaultWeightOnlyQuantConfig |
| t = time.time() |
| m = onnx.load(src) |
| cfg = DefaultWeightOnlyQuantConfig(block_size=32, is_symmetric=True, accuracy_level=4, |
| op_types_to_quantize=("MatMul", "Gather"), quant_axes=(("MatMul", 0), ("Gather", 1))) |
| q = MatMulNBitsQuantizer(m, algo_config=cfg) |
| q.process() |
| q.model.save_model_to_file(os.path.join(outdir, "model_int4.onnx"), use_external_data_format=False) |
| print("int4 %.0fs" % (time.time() - t)) |
| for f in sorted(os.listdir(outdir)): |
| print(f, os.path.getsize(os.path.join(outdir, f)) // 2**20, "MB") |
|
|
