File size: 4,179 Bytes
4dce182 e24bf15 4dce182 | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """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()
|