306 lines
9.5 KiB
Python
306 lines
9.5 KiB
Python
"""Export trained PyTorch checkpoints to ONNX for the inference SDK.
|
|
|
|
Usage:
|
|
uv run python -m tools.export_onnx \
|
|
--flow-ckpt data/models_v2 \
|
|
--scroll-ckpt data/scroll_models \
|
|
--output src/ai_mouse/assets/
|
|
|
|
Produces:
|
|
<output>/flow_model.onnx
|
|
<output>/scroll_decoder.onnx
|
|
<output>/click_dist.json
|
|
<output>/duration_dist.json
|
|
<output>/train_config.json
|
|
<output>/scroll_config.json
|
|
|
|
A PyTorch vs ONNX Runtime parity check runs at the end. If parity fails
|
|
the .onnx files are deleted to prevent shipping broken weights.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_ATOL = 1e-4
|
|
|
|
|
|
def export_flow_model(ckpt_dir: Path, out_dir: Path) -> Path:
|
|
"""Export TrajectoryFlowModel to ONNX.
|
|
|
|
Args:
|
|
ckpt_dir: directory with flow_model.pt and train_config.json.
|
|
out_dir: destination directory (created if missing).
|
|
|
|
Returns:
|
|
Path to the written flow_model.onnx.
|
|
"""
|
|
from tools.models import TrajectoryFlowModel
|
|
|
|
config_path = ckpt_dir / "train_config.json"
|
|
cfg = json.loads(config_path.read_text())
|
|
seq_len = int(cfg["seq_len"])
|
|
d_model = int(cfg["d_model"])
|
|
nhead = int(cfg["nhead"])
|
|
num_layers = int(cfg["num_layers"])
|
|
dim_feedforward = int(cfg["dim_feedforward"])
|
|
cond_dim = int(cfg.get("cond_dim", 3))
|
|
|
|
model = TrajectoryFlowModel(
|
|
seq_len=seq_len,
|
|
d_model=d_model,
|
|
nhead=nhead,
|
|
num_layers=num_layers,
|
|
dim_feedforward=dim_feedforward,
|
|
cond_dim=cond_dim,
|
|
dropout=0.0,
|
|
)
|
|
state = torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
|
|
model.load_state_dict(state)
|
|
model.eval()
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
out_path = out_dir / "flow_model.onnx"
|
|
|
|
dummy_x = torch.zeros(1, seq_len, 3, dtype=torch.float32)
|
|
dummy_t = torch.zeros(1, dtype=torch.float32)
|
|
dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32)
|
|
|
|
torch.onnx.export(
|
|
model,
|
|
(dummy_x, dummy_t, dummy_cond),
|
|
str(out_path),
|
|
input_names=["x_t", "t", "cond"],
|
|
output_names=["v"],
|
|
dynamic_axes={
|
|
"x_t": {0: "batch"},
|
|
"t": {0: "batch"},
|
|
"cond": {0: "batch"},
|
|
"v": {0: "batch"},
|
|
},
|
|
opset_version=17,
|
|
do_constant_folding=True,
|
|
dynamo=False,
|
|
)
|
|
logger.info("Wrote %s (%.1f MB)", out_path, out_path.stat().st_size / 1e6)
|
|
return out_path
|
|
|
|
|
|
class _ScrollDecoder(torch.nn.Module):
|
|
"""Wraps ScrollCVAE.decode for ONNX export.
|
|
|
|
The full ScrollCVAE is encoder+decoder; inference only needs decoder.
|
|
"""
|
|
|
|
def __init__(self, dec_h0, dec_gru, dec_out, seq_len: int, hidden: int):
|
|
super().__init__()
|
|
self.dec_h0 = dec_h0
|
|
self.dec_gru = dec_gru
|
|
self.dec_out = dec_out
|
|
self.seq_len = seq_len
|
|
self.hidden = hidden
|
|
|
|
def forward(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
|
|
b = z.shape[0]
|
|
zc = torch.cat([z, cond], dim=-1)
|
|
h0_flat = self.dec_h0(zc)
|
|
h0 = h0_flat.view(b, 2, self.hidden).permute(1, 0, 2).contiguous()
|
|
inp = zc.unsqueeze(1).expand(b, self.seq_len, -1)
|
|
out, _ = self.dec_gru(inp, h0)
|
|
return self.dec_out(out)
|
|
|
|
|
|
def export_scroll_decoder(ckpt_dir: Path, out_dir: Path) -> Path:
|
|
"""Export ScrollCVAE decoder to ONNX."""
|
|
from tools.scroll.models import ScrollCVAE
|
|
|
|
config_path = ckpt_dir / "scroll_config.json"
|
|
cfg = json.loads(config_path.read_text())
|
|
seq_len = int(cfg["seq_len"])
|
|
latent_dim = int(cfg["latent_dim"])
|
|
hidden = int(cfg["hidden"])
|
|
cond_dim = int(cfg["cond_dim"])
|
|
|
|
full = ScrollCVAE(
|
|
seq_len=seq_len, latent_dim=latent_dim, hidden=hidden, cond_dim=cond_dim
|
|
)
|
|
state = torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
|
|
full.load_state_dict(state)
|
|
full.eval()
|
|
|
|
decoder = _ScrollDecoder(
|
|
dec_h0=full.dec_h0,
|
|
dec_gru=full.dec_gru,
|
|
dec_out=full.dec_out,
|
|
seq_len=seq_len,
|
|
hidden=hidden,
|
|
)
|
|
decoder.eval()
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
out_path = out_dir / "scroll_decoder.onnx"
|
|
|
|
dummy_z = torch.zeros(1, latent_dim, dtype=torch.float32)
|
|
dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32)
|
|
|
|
torch.onnx.export(
|
|
decoder,
|
|
(dummy_z, dummy_cond),
|
|
str(out_path),
|
|
input_names=["z", "cond"],
|
|
output_names=["seq"],
|
|
dynamic_axes={
|
|
"z": {0: "batch"},
|
|
"cond": {0: "batch"},
|
|
"seq": {0: "batch"},
|
|
},
|
|
opset_version=17,
|
|
do_constant_folding=True,
|
|
dynamo=False,
|
|
)
|
|
logger.info("Wrote %s (%.1f KB)", out_path, out_path.stat().st_size / 1e3)
|
|
return out_path
|
|
|
|
|
|
def _check_flow_parity(ckpt_dir: Path, onnx_path: Path) -> None:
|
|
"""Verify ONNX flow model matches PyTorch output on random input."""
|
|
import onnxruntime as ort
|
|
from tools.models import TrajectoryFlowModel
|
|
|
|
cfg = json.loads((ckpt_dir / "train_config.json").read_text())
|
|
seq_len = int(cfg["seq_len"])
|
|
cond_dim = int(cfg.get("cond_dim", 3))
|
|
|
|
model = TrajectoryFlowModel(
|
|
seq_len=seq_len,
|
|
d_model=int(cfg["d_model"]),
|
|
nhead=int(cfg["nhead"]),
|
|
num_layers=int(cfg["num_layers"]),
|
|
dim_feedforward=int(cfg["dim_feedforward"]),
|
|
cond_dim=cond_dim,
|
|
dropout=0.0,
|
|
)
|
|
model.load_state_dict(
|
|
torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
|
|
)
|
|
model.eval()
|
|
|
|
torch.manual_seed(42)
|
|
np.random.seed(42)
|
|
x = torch.randn(2, seq_len, 3, dtype=torch.float32)
|
|
t = torch.tensor([0.0, 0.5], dtype=torch.float32)
|
|
cond = torch.randn(2, cond_dim, dtype=torch.float32)
|
|
|
|
with torch.no_grad():
|
|
torch_out = model(x, t, cond).numpy()
|
|
|
|
sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
|
|
ort_out = sess.run(
|
|
["v"],
|
|
{
|
|
"x_t": x.numpy(),
|
|
"t": t.numpy(),
|
|
"cond": cond.numpy(),
|
|
},
|
|
)[0]
|
|
|
|
if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
|
|
max_diff = float(np.abs(torch_out - ort_out).max())
|
|
raise RuntimeError(
|
|
f"Flow model ORT/PyTorch parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
|
|
)
|
|
logger.info("Flow model parity OK (atol=%.0e)", _ATOL)
|
|
|
|
|
|
def _check_scroll_parity(ckpt_dir: Path, onnx_path: Path) -> None:
|
|
"""Verify ONNX scroll decoder matches PyTorch decoder output."""
|
|
import onnxruntime as ort
|
|
from tools.scroll.models import ScrollCVAE
|
|
|
|
cfg = json.loads((ckpt_dir / "scroll_config.json").read_text())
|
|
seq_len = int(cfg["seq_len"])
|
|
latent_dim = int(cfg["latent_dim"])
|
|
cond_dim = int(cfg["cond_dim"])
|
|
|
|
full = ScrollCVAE(
|
|
seq_len=seq_len,
|
|
latent_dim=latent_dim,
|
|
hidden=int(cfg["hidden"]),
|
|
cond_dim=cond_dim,
|
|
)
|
|
full.load_state_dict(
|
|
torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
|
|
)
|
|
full.eval()
|
|
|
|
torch.manual_seed(7)
|
|
z = torch.randn(2, latent_dim, dtype=torch.float32)
|
|
cond = torch.randn(2, cond_dim, dtype=torch.float32)
|
|
|
|
with torch.no_grad():
|
|
torch_out = full.decode(z, cond).numpy()
|
|
|
|
sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
|
|
ort_out = sess.run(["seq"], {"z": z.numpy(), "cond": cond.numpy()})[0]
|
|
|
|
if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
|
|
max_diff = float(np.abs(torch_out - ort_out).max())
|
|
raise RuntimeError(
|
|
f"Scroll decoder parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
|
|
)
|
|
logger.info("Scroll decoder parity OK (atol=%.0e)", _ATOL)
|
|
|
|
|
|
def _copy_metadata(flow_dir: Path, scroll_dir: Path, out_dir: Path) -> None:
|
|
"""Copy JSON metadata files alongside the ONNX models."""
|
|
for name in ("click_dist.json", "duration_dist.json", "train_config.json"):
|
|
src = flow_dir / name
|
|
if not src.exists():
|
|
raise FileNotFoundError(f"Required metadata missing: {src}")
|
|
shutil.copy2(src, out_dir / name)
|
|
src = scroll_dir / "scroll_config.json"
|
|
if not src.exists():
|
|
raise FileNotFoundError(f"Required metadata missing: {src}")
|
|
shutil.copy2(src, out_dir / "scroll_config.json")
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
p = argparse.ArgumentParser(prog="export_onnx", description=__doc__.splitlines()[0])
|
|
p.add_argument("--flow-ckpt", type=Path, required=True)
|
|
p.add_argument("--scroll-ckpt", type=Path, required=True)
|
|
p.add_argument("--output", type=Path, required=True)
|
|
args = p.parse_args(argv)
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
|
|
flow_onnx = export_flow_model(args.flow_ckpt, args.output)
|
|
scroll_onnx = export_scroll_decoder(args.scroll_ckpt, args.output)
|
|
|
|
try:
|
|
_check_flow_parity(args.flow_ckpt, flow_onnx)
|
|
_check_scroll_parity(args.scroll_ckpt, scroll_onnx)
|
|
except RuntimeError as exc:
|
|
logger.error("Parity check failed: %s", exc)
|
|
flow_onnx.unlink(missing_ok=True)
|
|
scroll_onnx.unlink(missing_ok=True)
|
|
return 1
|
|
|
|
_copy_metadata(args.flow_ckpt, args.scroll_ckpt, args.output)
|
|
logger.info("Export complete: %s", args.output)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|