Spaces:
Sleeping
Sleeping
File size: 1,472 Bytes
fce5ca7 | 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 34 35 36 37 38 39 40 41 42 43 | """Create an optional INT8 ONNX variant; evaluate it before selecting for serving."""
import argparse
import json
import shutil
from pathlib import Path
from onnxruntime.quantization import quantize_dynamic, QuantType
from image_to_latex.artifacts import digest
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--source", type=Path, default=Path("artifacts/latex-ocr"))
p.add_argument("--output", type=Path, default=Path("artifacts/latex-ocr-int8"))
args = p.parse_args()
if args.output.exists():
p.error("Output must be new; original artifacts will not be overwritten")
args.output.mkdir(parents=True)
for name in ["encoder", "decoder"]:
quantize_dynamic(
str(args.source / f"{name}.onnx"),
str(args.output / f"{name}.int8.onnx"),
weight_type=QuantType.QInt8,
per_channel=True,
op_types_to_quantize=["MatMul", "Gemm"],
extra_options={"MatMulConstBOnly": True},
)
for name in ["image_resizer.onnx", "tokenizer.json"]:
shutil.copy2(args.source / name, args.output / name)
metadata = {
p.name: dict(bytes=p.stat().st_size, sha256=digest(p))
for p in args.output.iterdir()
if p.is_file()
}
(args.output / "quantization.json").write_text(
json.dumps(metadata, indent=2) + "\n"
)
print(json.dumps(metadata, indent=2))
if __name__ == "__main__":
main()
|