#!/usr/bin/env python3 """Generate optimized ECSeg ONNX variants from an FP32 baseline. Every variant PRESERVES the public graph I/O contract (input `images` float32 [1,3,640,640]; outputs labels int64 [1,300], boxes/scores/masks float32) because the app's curated-import verification checks tensor dtype+shape against a pinned contract. A variant that changed an output dtype would fail that gate, so we never let optimization touch the boundary tensors. Subcommands (each writes /..onnx and runs onnx.checker): graphopt onnxslim constant-folding + lossless simplification (runtime-agnostic; no ORT-CPU-specific fusions baked in, so it stays portable to ort-web WASM) fp16 onnxconverter-common float16 weights, keep_io_types=True (Cast at boundaries), with a conservative op_block_list for ops whose fp16 kernels are risky on WASM int8-dynamic weight-only dynamic INT8 (per-channel) on MatMul/Gemm int8-static static QDQ INT8 with real calibration images (Conv+MatMul), optional exclude list Usage: python optimize.py graphopt IN.onnx OUT_DIR python optimize.py fp16 IN.onnx OUT_DIR python optimize.py int8-dynamic IN.onnx OUT_DIR python optimize.py int8-static IN.onnx OUT_DIR --calib-dir DIR [--calib-limit N] [--exclude-mask-head] """ from __future__ import annotations import argparse import os import sys import onnx HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) import ecseg_common as ec # noqa: E402 def _check(path: str) -> None: model = onnx.load(path, load_external_data=False) onnx.checker.check_model(model) print(f" onnx.checker OK: {path} ({os.path.getsize(path)/2**20:.2f} MiB)") def _stem(in_path: str) -> str: base = os.path.basename(in_path) return base[:-5] if base.endswith(".onnx") else base # --------------------------------------------------------------------------- graphopt -------- def graphopt(in_path: str, out_dir: str) -> str: import onnxslim out = os.path.join(out_dir, f"{_stem(in_path)}.graphopt.onnx") # onnxslim performs constant folding, dead-node elimination and safe fusions while keeping the # graph runtime-agnostic (unlike ORT's offline optimizer, which can bake in CPU-specific layout # fusions that don't transfer to ort-web WASM). model = onnxslim.slim(in_path) onnx.save(model, out) _check(out) return out # ------------------------------------------------------------------------------ fp16 --------- # Ops we keep in fp32 even inside an fp16 graph: boundary-sensitive / sampling / reduction ops # whose fp16 kernels are either absent or numerically fragile on the ort-web WASM CPU EP. FP16_OP_BLOCK_LIST = [ "GridSample", # deformable mask sampling — precision-critical, fp16 kernel support uncertain "Resize", "LayerNormalization", "Softmax", "ReduceSum", "ReduceMean", "ReduceMax", "Erf", "TopK", "GatherElements", "Einsum", ] def fp16(in_path: str, out_dir: str, hi_range: bool = False) -> str: from onnxconverter_common import float16 suffix = "fp16-hi" if hi_range else "fp16" out = os.path.join(out_dir, f"{_stem(in_path)}.{suffix}.onnx") model = onnx.load(in_path) kwargs = dict( keep_io_types=True, # inputs/outputs stay float32 -> contract preserved disable_shape_infer=False, op_block_list=FP16_OP_BLOCK_LIST, ) if hi_range: # Use fp16's true representable range instead of the library's conservative ±1e4 clamp, so large # mask logits are not clipped. min_positive_val = smallest fp16 normal (2**-14). kwargs["max_finite_val"] = 65504.0 kwargs["min_positive_val"] = 5.9604645e-08 converted = float16.convert_float_to_float16(model, **kwargs) onnx.save(converted, out) _check(out) return out # ------------------------------------------------------------------------ int8-dynamic ------- def int8_dynamic(in_path: str, out_dir: str) -> str: from onnxruntime.quantization import QuantType, quantize_dynamic from onnxruntime.quantization.shape_inference import quant_pre_process stem = _stem(in_path) pre = os.path.join(out_dir, f"{stem}.preproc.onnx") quant_pre_process(in_path, pre, skip_symbolic_shape=False) out = os.path.join(out_dir, f"{stem}.int8-dynamic.onnx") quantize_dynamic( pre, out, weight_type=QuantType.QInt8, per_channel=True, # Conv dynamic quantization is not well supported; restrict to the matmul/gemm family. op_types_to_quantize=["MatMul", "Gemm"], ) os.remove(pre) _check(out) return out # ------------------------------------------------------------------------- int8-static ------- class _CalibReader: """onnxruntime CalibrationDataReader over app-faithful preprocessed images.""" def __init__(self, files, input_name="images"): self.files = files self.input_name = input_name self._it = iter(files) def get_next(self): path = next(self._it, None) if path is None: return None return {self.input_name: ec.preprocess_file(path)} def rewind(self): self._it = iter(self.files) # Mask-head node-name substrings to optionally keep in fp32 (the segmentation branch is the most # quantization-sensitive part). Discovered empirically from the graph; matched case-insensitively. MASK_HEAD_HINTS = ["mask", "seg", "pixel_decoder", "GridSample"] def int8_static( in_path: str, out_dir: str, calib_dir, calib_limit=None, exclude_mask_head=False, ) -> str: from onnxruntime.quantization import ( CalibrationMethod, QuantFormat, QuantType, quantize_static, ) from onnxruntime.quantization.shape_inference import quant_pre_process stem = _stem(in_path) suffix = "int8-static-selective" if exclude_mask_head else "int8-static" pre = os.path.join(out_dir, f"{stem}.preproc.onnx") quant_pre_process(in_path, pre, skip_symbolic_shape=False) files = ec.list_images(calib_dir if isinstance(calib_dir, list) else [calib_dir], limit=calib_limit) if not files: raise SystemExit(f"No calibration images found under {calib_dir}") print(f" calibration images: {len(files)}") nodes_to_exclude = [] if exclude_mask_head: model = onnx.load(pre, load_external_data=False) for node in model.graph.node: name = (node.name or "").lower() if any(h.lower() in name for h in MASK_HEAD_HINTS): nodes_to_exclude.append(node.name) print(f" excluding {len(nodes_to_exclude)} mask-head nodes from quantization") out = os.path.join(out_dir, f"{stem}.{suffix}.onnx") quantize_static( pre, out, _CalibReader(files), quant_format=QuantFormat.QDQ, per_channel=True, weight_type=QuantType.QInt8, activation_type=QuantType.QUInt8, calibrate_method=CalibrationMethod.MinMax, nodes_to_exclude=nodes_to_exclude, extra_options={"AddQDQPairToWeight": False}, ) os.remove(pre) _check(out) return out # -------------------------------------------------------------------------------- cli -------- def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="cmd", required=True) for name in ("graphopt", "fp16", "int8-dynamic"): p = sub.add_parser(name) p.add_argument("in_path") p.add_argument("out_dir") if name == "fp16": p.add_argument("--hi-range", action="store_true", help="use fp16's full ±65504 range instead of the default ±1e4 clamp") ps = sub.add_parser("int8-static") ps.add_argument("in_path") ps.add_argument("out_dir") ps.add_argument("--calib-dir", action="append", required=True) ps.add_argument("--calib-limit", type=int, default=None) ps.add_argument("--exclude-mask-head", action="store_true") args = ap.parse_args() os.makedirs(args.out_dir, exist_ok=True) if args.cmd == "graphopt": graphopt(args.in_path, args.out_dir) elif args.cmd == "fp16": fp16(args.in_path, args.out_dir, hi_range=getattr(args, "hi_range", False)) elif args.cmd == "int8-dynamic": int8_dynamic(args.in_path, args.out_dir) elif args.cmd == "int8-static": int8_static( args.in_path, args.out_dir, args.calib_dir, args.calib_limit, args.exclude_mask_head ) return 0 if __name__ == "__main__": sys.exit(main())