96 lines
2.6 KiB
Python
96 lines
2.6 KiB
Python
"""Export trained PyTorch checkpoints to ONNX for the inference SDK.
|
|
|
|
Usage:
|
|
uv run python tools/export_onnx.py \
|
|
--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
|