alexvoss's picture
Publish ECSeg optimization study and experimental FP16 artifacts
9496f98 verified
|
Raw
History Blame Contribute Delete
6.74 kB

ECSeg ONNX optimization pipeline

Reproducible tooling that answers one question: should AnnotateIt replace the FP32 ECSeg-S / ECSeg-M checkpoints with an optimized variant? It downloads the pinned FP32 baselines, builds candidate variants (lossless graph-opt, FP16, dynamic INT8, static INT8/QDQ), checks numerical equivalence against FP32, and β€” decisively β€” measures every variant in the target runtime (onnxruntime-web 1.24.3, CPU/WASM) via a real browser, not Python.

The full write-up lives in docs/model-optimization/ecseg-quantization-report.md.

Nothing here is an app dependency. The Python tools run in a throwaway venv (see requirements.txt); the browser benchmark uses the repo's own node_modules/onnxruntime-web. No generated .onnx is committed β€” they are large and rebuildable.

Layout

File Purpose
inspect_onnx.py opset / node & initializer counts / weight dtypes / external-data / I/O contract / op histogram
optimize.py build variants: graphopt, fp16, int8-dynamic, int8-static (all preserve the public I/O contract)
ecseg_common.py app-faithful preprocessing (stretch-640 bilinear + ImageNet) + edgecrafter-seg post-processing
correctness.py candidate-vs-FP32 raw-output error + end-to-end instance/box/mask-IoU on identical inputs
aggregate_report.py consolidate report JSONs β†’ combined.json + CSVs + markdown tables
bench/harness.{html,mjs} in-browser ORT harness, configured exactly like the app's ECSeg session
bench/bench_browser.mjs Playwright driver: serves harness+ORT+models to Chrome/WebKit, measures latency

Prerequisites

# Python 3.12 (onnxruntime has no 3.14 wheels); create the isolated venv:
python3.12 -m venv .venv && source .venv/bin/activate
pip install -r scripts/model-optimization/requirements.txt

# Node deps (playwright, onnxruntime-web) are already in the repo's node_modules.
# The browser benchmark drives your system Google Chrome (Playwright channel 'chrome').

Reproduce

Set a work dir for artifacts (kept OUT of git):

export WORK=/tmp/ecseg-opt && mkdir -p $WORK/{baseline,variants,bench-inputs,coco/val2017,reports}

1. Baseline (download + verify pinned SHA-256)

curl -sL -o $WORK/baseline/ecseg-s.onnx \
  https://huggingface.co/AnnotateIt/edgecrafter-ecseg-s-onnx/resolve/a71d0ee5e27b1083d40cd5ced647139a7ab8d533/model.onnx
curl -sL -o $WORK/baseline/ecseg-m.onnx \
  https://huggingface.co/AnnotateIt/edgecrafter-ecseg-m-onnx/resolve/56812fa474791c8db90740141e359bc767e0dc22/model.onnx
shasum -a 256 $WORK/baseline/ecseg-s.onnx  # fd30a61cee9b798259e89e5b58cbe57a7c65da2b3067f33ccf02b923402522da
shasum -a 256 $WORK/baseline/ecseg-m.onnx  # ded8147572ecec27fe2b384cfac5e5cc38738c9b3be773d5ccbee3c261820474

python scripts/model-optimization/inspect_onnx.py $WORK/baseline/ecseg-s.onnx --json $WORK/reports/inspect-ecseg-s-fp32.json

2. Calibration / evaluation images

The report calibrated static INT8 on a 50-image COCO val2017 subset and evaluated agreement on the repo's held-out datasets/ images (disjoint from calibration). To refetch the COCO subset, use the id list the report was produced with ($WORK/coco/ids.txt) and pull from http://images.cocodataset.org/val2017/<id>.jpg.

3. Build variants (<model> = ecseg-s or ecseg-m)

V=$WORK/variants
python scripts/model-optimization/optimize.py graphopt     $WORK/baseline/<model>.onnx $V
python scripts/model-optimization/optimize.py fp16         $WORK/baseline/<model>.onnx $V
python scripts/model-optimization/optimize.py int8-dynamic $WORK/baseline/<model>.onnx $V
python scripts/model-optimization/optimize.py int8-static  $WORK/baseline/<model>.onnx $V --calib-dir $WORK/coco/val2017
python scripts/model-optimization/optimize.py int8-static  $WORK/baseline/<model>.onnx $V --calib-dir $WORK/coco/val2017 --exclude-mask-head
cp $WORK/baseline/<model>.onnx $V/<model>.fp32.onnx

4. Numerical equivalence vs FP32

python scripts/model-optimization/correctness.py \
  --baseline $WORK/baseline/ecseg-s.onnx \
  --candidate $WORK/variants/ecseg-s.fp16.onnx \
  --images datasets --json $WORK/reports/corr-ecseg-s-fp16.json

5. Browser latency in the target runtime (the decisive measurement)

# Export byte-identical preprocessed inputs (so browser feeds match the Python correctness run):
python - <<'PY'
import sys, numpy as np; sys.path.insert(0,'scripts/model-optimization'); import ecseg_common as ec, os
W=os.environ['WORK']
for n in ['000000000139','000000000785','000000002149']:
    ec.preprocess_file(f'{W}/coco/val2017/{n}.jpg').astype(np.float32).tofile(f'{W}/bench-inputs/{n}.f32')
PY

# wasm-mt (cross-origin isolated, threaded pool) must run HEADED β€” headless Chrome will not bring up
# ORT's WASM worker pool (session.create hangs). wasm-st can run headless (BENCH_HEADLESS=1).
node scripts/model-optimization/bench/bench_browser.mjs \
  --models-dir $WORK/variants --input-dir $WORK/bench-inputs \
  --models ecseg-s.fp32.onnx,ecseg-s.fp16.onnx,ecseg-s.int8-dynamic.onnx,ecseg-s.int8-static.onnx,ecseg-s.int8-static-selective.onnx,ecseg-s.graphopt.onnx \
  --configs wasm-mt,wasm-st --engine chrome --repeats 15 --warmups 2 \
  --out $WORK/reports/results-s.json

# Optional second engine (Safari/WebKit) for cross-browser confirmation:
#   --engine webkit

6. Consolidate

python scripts/model-optimization/aggregate_report.py --reports-dir $WORK/reports --variants-dir $WORK/variants

Runtime facts the tooling encodes (from the app source)

  • Storage is uncompressed β€” the offline model store keeps the raw ONNX bytes in IndexedDB (offline-model-store.ts). gzip/Brotli on the wire does not shrink the stored artifact; disk bytes β‰ˆ IndexedDB bytes. That is why the size tables use on-disk .onnx size.
  • The I/O contract is verified strictly at curated import (custom-model-curated-flow.ts verifyTensorAgainstSpec): tensor dtype and shape must match images f32[1,3,640,640] β†’ labels i64[1,300], boxes/scores f32, masks f32[1,300,160,160]. Every variant preserves this (FP16 keeps IO in float32 via boundary casts; INT8 keeps IO float).
  • ECSeg is pinned to the CPU/WASM EP (custom-checkpoints.ts) β€” WebGPU fails this graph on ort-web 1.24.3 (the graph uses GridSample). The benchmark therefore measures CPU/WASM only.
  • Session options mirrored by the harness: executionProviders:['cpu'], graphOptimizationLevel:'all', logSeverityLevel:3, executionMode sequential (1 thread) / parallel (pool), env.wasm.simd=true.