| """Export trained components to ONNX for CPU inference with ONNX Runtime, and check parity. |
| |
| python -m tinydoc.export --weights runs --out models |
| """ |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from .features import MAX_CHARS, NUM_FEATS |
| from .models import build, count_params |
|
|
| OPSET = 17 |
|
|
|
|
| def sample_inputs(name): |
| g = torch.Generator().manual_seed(0) |
| if name == "detector": |
| return (torch.rand(1, 3, 640, 480, generator=g),), ["image"], ["kernel_logits"], {"image": {0: "batch", 2: "height", 3: "width"}} |
| if name == "recognizer": |
| x = torch.rand(2, 1, 32, 320, generator=g) |
| mask = torch.zeros(2, 80, dtype=torch.bool) |
| mask[1, 60:] = True |
| return (x, mask), ["image", "step_mask"], ["logits"], {"image": {0: "batch", 3: "width"}, "step_mask": {0: "batch", 1: "steps"}} |
| if name == "orientation": |
| return (torch.rand(4, 1, 192, 192, generator=g),), ["patches"], ["logits"], {"patches": {0: "batch"}} |
| n = 37 |
| return ((torch.randint(1, 100, (1, n, MAX_CHARS), generator=g), torch.rand(1, n, NUM_FEATS, generator=g), |
| torch.rand(1, n, 4, generator=g) * 30, torch.ones(1, n)), |
| ["char_ids", "feats", "rel_boxes", "valid"], ["field_logits", "rel_logits", "doc_logits"], |
| {"char_ids": {0: "batch", 1: "segments"}, "feats": {0: "batch", 1: "segments"}, |
| "rel_boxes": {0: "batch", 1: "segments"}, "valid": {0: "batch", 1: "segments"}}) |
|
|
|
|
| def alt_inputs(name): |
| """Different shapes from the export sample, to prove the dynamic axes really are dynamic.""" |
| g = torch.Generator().manual_seed(1) |
| if name == "detector": |
| return (torch.rand(1, 3, 320, 704, generator=g),) |
| if name == "recognizer": |
| return torch.rand(3, 1, 32, 512, generator=g), torch.zeros(3, 128, dtype=torch.bool) |
| if name == "orientation": |
| return (torch.rand(7, 1, 192, 192, generator=g),) |
| n = 120 |
| return (torch.randint(1, 100, (1, n, MAX_CHARS), generator=g), torch.rand(1, n, NUM_FEATS, generator=g), |
| torch.rand(1, n, 4, generator=g) * 30, torch.ones(1, n)) |
|
|
|
|
| def export_one(name, weights, out_dir, random_init=False): |
| import onnxruntime as ort |
|
|
| ck = None if random_init else torch.load(Path(weights) / name / "best.pt", map_location="cpu", weights_only=False) |
| model = build(name, **((ck or {}).get("arch") or {})) |
| if ck: |
| model.load_state_dict(ck["model"]) |
| model.eval() |
| args, in_names, out_names, dyn = sample_inputs(name) |
| path = Path(out_dir) / f"{name}.onnx" |
| torch.onnx.export(model, args, str(path), input_names=in_names, output_names=out_names, |
| dynamic_axes=dyn, opset_version=OPSET, dynamo=False) |
| import onnx |
| proto = onnx.load(str(path)) |
| meta = proto.metadata_props.add() |
| meta.key, meta.value = "trained_parameters", str(count_params(model)) |
| onnx.save(proto, str(path)) |
| sess = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) |
| worst = 0.0 |
| for inputs in (args, alt_inputs(name)): |
| with torch.no_grad(): |
| ref = model(*inputs) |
| ref = ref if isinstance(ref, tuple) else (ref,) |
| got = sess.run(None, {n: t.numpy() for n, t in zip(in_names, inputs)}) |
| worst = max(worst, max(float(np.abs(r.numpy() - g).max()) for r, g in zip(ref, got))) |
| size = path.stat().st_size / 1e6 |
| print(f" {name:12s} params={count_params(model):>9,d} {size:5.1f} MB max |torch - onnx| = {worst:.2e}") |
| return worst |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--weights", default="runs") |
| ap.add_argument("--out", default="models") |
| ap.add_argument("--random-init", action="store_true", help="export untrained weights (pipeline testing)") |
| args = ap.parse_args() |
| Path(args.out).mkdir(parents=True, exist_ok=True) |
| worst = max(export_one(n, args.weights, args.out, args.random_init) for n in ("detector", "recognizer", "orientation", "extractor")) |
| if worst > 1e-3: |
| raise SystemExit(f"ONNX parity check failed (max diff {worst:.2e})") |
| print(f"exported to {args.out}/ (parity ok)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
