# ai_mouse 库化与 ONNX Runtime 迁移 — 设计文档 **日期**: 2026-05-11 **状态**: Draft (待用户审阅) **目标版本**: ai_mouse 0.2.0 ## 背景与目标 当前 `ai_mouse` 仓库既是应用又是包:推理代码、训练代码、Web UI、采集、评估、数据适配器都混在 `ai_mouse/` 包目录下,运行时强依赖 `torch`、`fastapi`、`scipy`、`matplotlib`。这与"被其他项目直接引用的库"定位不符——任何只想调用 `generate()` 的下游必须装一整套 GB 级依赖。 本次重构的目标: 1. **物理边界**: 推理库代码与开发期代码完全分开,推理库不依赖 torch 2. **ONNX 化**: 推理走 ONNX Runtime,预训练权重以 ONNX 格式打进 wheel 3. **可消费**: 下游 `pip install git+...` 即可 `from ai_mouse import generate` 4. **API 优化**: 同时提供函数式入口(向后兼容)和类入口(性能场景),底层共享 ORT session 非目标(显式排除): - 不发布到 PyPI(本阶段仅 git URL 安装) - 不重写训练代码以脱离 torch - 不做 ONNX Runtime Training - 不为视觉真实性设硬阈值断言 ## 约束(由用户决策确认) | 维度 | 决策 | |---|---| | 使用场景 | 纯推理 SDK,零 torch 运行时依赖 | | 权重交付 | 打进 wheel(`importlib.resources`) | | 开发期代码位置 | 仓库顶层 `tools/`,不进 wheel | | 公开 API 形态 | 类 + 函数双轨;函数内部走进程级 lru_cache 单例 | | 发布渠道 | git URL 直装,不发 PyPI | ## 设计 ### §1 公开 API #### 顶层入口 (`src/ai_mouse/__init__.py`) ```python from ai_mouse import ( MouseModel, # 类: 常驻 InferenceSession ScrollModel, generate, # 便利函数: 进程级单例缓存 generate_scroll, ) from ai_mouse import errors # ModelLoadError, GenerationError __version__ = "0.2.0" ``` #### `MouseModel` ```python class MouseModel: def __init__( self, model_path: str | Path | None = None, # None ⇒ wheel 内置权重 providers: Sequence[str] | None = None, # None ⇒ ORT 默认顺序 (CPU) seed: int | None = None, # None ⇒ 每次调用随机 ) -> None: ... def generate( self, start: tuple[int, int], end: tuple[int, int], n_points: int = 64, speed: float | None = None, click: bool = True, # True ⇒ 末尾追加 mouse-down/-up seed: int | None = None, # 覆盖实例 seed ) -> list[tuple[int, int, int]]: ... def sample_click_duration_ms(self, seed: int | None = None) -> int: ... def close(self) -> None: ... def __enter__(self) -> "MouseModel": ... def __exit__(self, *exc) -> None: ... ``` #### `ScrollModel`(对称) ```python class ScrollModel: def __init__(self, model_path=None, providers=None, seed=None): ... def generate( self, start_scroll_y: int, target_scroll_y: int, mode: Literal["target", "fast", "precise"] = "target", viewport_height: int = 800, seed: int | None = None, ) -> list[dict]: ... def close(self) -> None: ... def __enter__(self) -> "ScrollModel": ... def __exit__(self, *exc) -> None: ... ``` #### 便利函数 ```python def generate(start, end, *, n_points=64, speed=None, click=True, seed=None, model_path=None, providers=None) -> list[tuple[int, int, int]]: ... def generate_scroll(start_scroll_y, target_scroll_y, *, mode="target", viewport_height=800, seed=None, model_path=None, providers=None) -> list[dict]: ... ``` - 内部以 `(model_path or "__bundled__", tuple(providers or ()))` 为 key,`functools.lru_cache(maxsize=4)` 缓存模型实例 - 多次调用同参数 ⇒ 复用同一 `InferenceSession`,不重复加载 - `maxsize=4` 的选择: 典型场景只用 bundled 权重 + 默认 provider(size=1 够用);保留余量给"自定义权重 + CPU/GPU 双 provider"等极少数场景。再大没意义,大模型对象常驻内存反而浪费 #### 与旧 API 的关键差异 | 旧 | 新 | 原因 | |---|---|---| | `model_dir: str` | `model_path: str \| Path` | 接受 Path,语义更明确 | | `config: GenerateConfig` 参数 | 移除 | 内部物;可调项直接 kwarg | | 无显式 click 开关 | `click: bool = True` | 默认行为不变,但允许纯轨迹 | | 隐式 `viewport_norm` | `viewport_height: int = 800` | 暴露给调用者 | | 每次调用可能重载模型 | 类持久 / 函数走 lru_cache | SDK 不能"每次几百 ms 起步" | #### 线程安全 ORT `InferenceSession.run` 本身线程安全。`MouseModel`/`ScrollModel` 共享一个 session,可在多线程并发 `.generate()`。 `seed=None` 时使用 `np.random.default_rng()` 局部实例,不污染全局 numpy 状态。 #### 错误层级 ```python class AiMouseError(Exception): ... class ModelLoadError(AiMouseError): ... # 权重缺失 / shape 不符 / providers 不可用 class GenerationError(AiMouseError): ... # 推理时数值异常 (NaN 等) ``` ### §2 ONNX 导出方案 #### 工具脚本 `tools/export_onnx.py` 一次性脚本(训练完成后跑),产出可烧进 wheel 的 `.onnx`。库代码不依赖 torch;只有这个脚本依赖。 ```bash uv run python tools/export_onnx.py \ --flow-ckpt data/models_v2 \ --scroll-ckpt data/scroll_models \ --output src/ai_mouse/assets/ ``` 行为: - 读 `train_config.json` ⇒ 用一致超参实例化 `TrajectoryFlowModel` - `torch.load(flow_model.pt)` ⇒ `model.eval()` ⇒ 导出为 `flow_model.onnx` - 同样导出 scroll **decoder** 部分 ⇒ `scroll_decoder.onnx`(encoder 仅训练用) - 复制 `click_dist.json`、`duration_dist.json`、`train_config.json`、`scroll_config.json` 到 `assets/` #### 轨迹模型导出 ```python torch.onnx.export( model, args=( torch.zeros(1, seq_len, 3), # x_t torch.zeros(1), # t torch.zeros(1, 3), # cond ), f="flow_model.onnx", 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, ) ``` 关键决策: - **`seq_len=64` 静态**: 位置嵌入是固定 shape 可学习参数;`n_points` API 若 ≠ 64,在后处理阶段线性插值重采样,**不进 ONNX** - **Batch 动态**: 支持下游批量生成 - **Opset 17**: 支持 SDPA。如有问题脚本回落到 14 - **常量折叠**: 位置嵌入、time/cond MLP 权重可折掉 #### Euler ODE 循环留在 Python 10 步 Euler 在 Python/NumPy 中循环,每步一次 `session.run()`: ```python x = rng.standard_normal((1, SEQ_LEN, 3), dtype=np.float32) dt = 1.0 / N_STEPS for i in range(N_STEPS): t = np.full((1,), (i + 0.5) * dt, dtype=np.float32) v = self.session.run(["v"], {"x_t": x, "t": t, "cond": cond_np})[0] x = x + dt * v ``` 权衡:10 次 ORT 调用开销 ≈ 1-3 ms;不引入 ONNX `Loop` op 复杂度;单文件权重;调试容易。未来如需性能可升级单图 Loop 版,本次不做。 #### 滚轮模型导出(只导 decoder) ```python class ScrollDecoder(nn.Module): """Wrap ScrollCVAE.decode for ONNX export.""" def forward(self, z, cond): ... torch.onnx.export( ScrollDecoder(scroll_cvae), args=(torch.zeros(1, LATENT_DIM), torch.zeros(1, 7)), f="scroll_decoder.onnx", input_names=["z", "cond"], output_names=["seq"], dynamic_axes={"z": {0: "b"}, "cond": {0: "b"}, "seq": {0: "b"}}, opset_version=17, ) ``` `z` 在 Python 端 `rng.standard_normal((1, LATENT_DIM))` 采样;ONNX 内不放采样算子。 #### 体积估算 | 文件 | 估算大小 | |---|---| | `flow_model.onnx` | ~2-3 MB (4 层 Transformer, d_model=128, ~600K 参数 FP32) | | `scroll_decoder.onnx` | <300 KB | | JSON 元数据 (3-4 个) | <20 KB | | **wheel 总增量** | **~3 MB** | 未来可选 `--quantize int8` 压到 ~1 MB,本次不做(YAGNI)。 #### 兼容性自检(导出脚本内置) 末尾跑烟雾测试: 1. ORT 加载刚导出的 `.onnx` 2. 与 PyTorch 模型同输入做前向 3. `np.allclose(torch_out, ort_out, atol=1e-4)` 必须通过 4. 失败 ⇒ 不写入 `assets/`,保留 `.pt` ### §3 后处理 / 坐标变换的纯 NumPy 重写 `coord.py`、`_gaussian_smooth`、`_sample_duration`、`utils.resample_arc`、所有空间/时间后处理已经是 numpy。需要替换的只有四处: | 现在用 | 替换为 | |---|---| | `torch.randn(...)` | `rng.standard_normal((1, T, 3), dtype=np.float32)` | | `model(x_t, t, cond)` | `session.run(["v"], {...})[0]` | | `torch.load(...) / .eval()` | `onnxruntime.InferenceSession(asset_path, providers=...)` | | `scipy.stats.truncnorm.rvs(...)` | 纯 numpy 拒绝采样(见下) | #### 库内部文件分布 ``` src/ai_mouse/ ├── __init__.py # 公开 API 拼装 ├── mouse.py # MouseModel: session + Euler 循环 + 后处理 + click ├── scroll.py # ScrollModel: decoder run + 量化 + 时间戳 ├── _coord.py # encode_trajectory, decode_trajectory ├── _postprocess.py # 平滑、单调性、重采样、截断正态、duration 采样 ... ├── _assets.py # importlib.resources 加载 onnx + json ├── errors.py # AiMouseError, ModelLoadError, GenerationError ├── py.typed # 空文件 └── assets/ ├── flow_model.onnx ├── scroll_decoder.onnx ├── click_dist.json ├── duration_dist.json ├── train_config.json └── scroll_config.json ``` 下划线开头的模块是私有实现,**不在 `__all__`**,版本间可自由破坏。 **`tools/` 与 `src/ai_mouse/` 的依赖方向**: `tools/` 是仓库内部的开发期代码,允许 `from ai_mouse._coord import ...` 反向引用库私有模块(它们在同一 repo 共同演化,只是 `tools/` 不进 wheel)。但 `src/ai_mouse/` **不允许** `import tools.*` —— wheel 必须自包含。CI 的 library job 会通过"无 torch 环境"间接守住这条规则。 #### `_truncnorm_sample`(替代 scipy) ```python def _truncnorm_sample( mu: float, sigma: float, low: float, high: float, rng: np.random.Generator, max_tries: int = 32, ) -> float: for _ in range(max_tries): v = rng.normal(mu, sigma) if low <= v <= high: return float(v) # 极少触发的兜底 return float(np.clip(rng.normal(mu, sigma), low, high)) ``` 默认参数 `mu=80, sigma=30, low=20, high=300`:接受率 ≈ 97%,无可测开销。 #### `_postprocess.py` 公开函数清单 ```python def snap_endpoints(forward, lateral, seq_len, n_snap=6) -> tuple[np.ndarray, np.ndarray]: ... def smooth_start(forward, lateral, n=4) -> tuple[np.ndarray, np.ndarray]: ... def enforce_forward_monotonic(forward) -> np.ndarray: ... def gaussian_smooth(x, sigma=1.0) -> np.ndarray: ... def build_timestamps(log_dt, total_duration_ms, dt_clip=(2.0, 150.0)) -> np.ndarray: ... def sample_duration(duration_dist, dist, rng) -> float: ... def truncnorm_sample(mu, sigma, low, high, rng) -> float: ... def resample_arc(pixels, n_points) -> np.ndarray: ... ``` 每个函数纯函数,显式接收 `rng`,便于单测。 #### Scroll 后处理改动 `scroll/generator.py` 已完全 numpy,迁移只换两行: - `torch.randn(1, latent_dim)` ⇒ `rng.standard_normal((1, latent_dim), dtype=np.float32)` - `model.decode(z, cond_t)` ⇒ `session.run(["seq"], {"z": z, "cond": cond})[0]` 其余 softmax 归一化、量化、时间戳构建照搬。 ### §4 仓库迁移 #### 终态布局 ``` ai_mouse/ (repo root) ├── src/ │ └── ai_mouse/ (唯一进 wheel 的目录) │ ├── __init__.py │ ├── mouse.py │ ├── scroll.py │ ├── _coord.py │ ├── _postprocess.py │ ├── _assets.py │ ├── errors.py │ ├── py.typed │ └── assets/ │ ├── flow_model.onnx │ ├── scroll_decoder.onnx │ ├── click_dist.json │ ├── duration_dist.json │ ├── train_config.json │ └── scroll_config.json ├── tools/ (开发期脚本,不进 wheel) │ ├── __init__.py │ ├── __main__.py (CLI dispatch, 从 ai_mouse/__main__.py 移来) │ ├── train.py │ ├── serve.py (从 ./main.py 移来) │ ├── export_onnx.py ★ 新 │ ├── trainer.py │ ├── models.py (TrajectoryFlowModel — torch) │ ├── collector.py │ ├── config.py (TrainConfig 等训练侧 dataclass) │ ├── server/ │ ├── eval/ │ ├── data_adapters/ │ └── scroll/ │ ├── trainer.py │ ├── models.py │ └── collector.py ├── data/ (不变) ├── docs/ (不变) ├── static/ (不变,tools/serve.py 引用) ├── tests/ │ ├── unit/ (库测试,只需 numpy + ort) │ │ ├── conftest.py │ │ ├── test_mouse.py │ │ ├── test_scroll.py │ │ ├── test_coord.py │ │ ├── test_postprocess.py │ │ ├── test_assets.py │ │ ├── test_errors.py │ │ └── data/ │ │ ├── golden_mouse.npz │ │ ├── golden_scroll.npz │ │ └── tiny_flow.onnx │ └── tools/ (训练/服务器测试,需 torch+fastapi+...) │ ├── conftest.py │ ├── test_trainer.py │ ├── test_server.py │ ├── test_eval_metrics.py │ ├── test_balabit_adapter.py │ ├── test_export_onnx.py │ └── test_scroll_trainer.py ├── examples/ │ └── quickstart.py ├── pyproject.toml (runtime: numpy+ort;[dev]: torch+fastapi+...) ├── README.md (SDK 视角重写) ├── CHANGELOG.md (新建,以 0.2.0 起锚) ├── CLAUDE.md (更新) └── uv.lock ``` 被**删除**: - `ai_mouse/generator.py` ⇒ 被 `src/ai_mouse/mouse.py` 取代 - `ai_mouse/scroll/generator.py` ⇒ 被 `src/ai_mouse/scroll.py` 取代 - `ai_mouse/utils.py`(`resample_arc` 移入 `_postprocess.py`) - `JointCVAE`(legacy) - `_BUNDLED_MODELS_DIR = Path(__file__).parent.parent / "data" / ...` 这类源码树相对路径 #### 迁移分阶段 **阶段 1 — 物理移动,训练侧继续工作** - `git mv` 训练 / 服务器 / 采集 / 评估 / 数据适配器进 `tools/` - 修订 `tools/` 内部 import 为绝对 `from tools.X` - `main.py` ⇒ `tools/serve.py` - `tests/test_*` 按归属分进 `tests/unit/` 或 `tests/tools/` - 验收: `python -m tools train ...`、`python -m tools serve` 可跑;`from ai_mouse import generate` 仍能用(老 torch 实现) **阶段 2 — src-layout 切换 + pyproject 更新** - `git mv ai_mouse src/ai_mouse` - pyproject 切到 hatchling,运行时依赖收紧到 `numpy + onnxruntime`;`[dev]` 群组收纳 torch+fastapi+scipy+matplotlib+pytest+... - 验收: `uv build` 产出 wheel,内容仅 `src/ai_mouse` **阶段 3 — ONNX 导出脚本 + 权重塞进 assets/** - 写 `tools/export_onnx.py` - 运行导出,产出 `.onnx`+元数据 commit 进库 - 验收: `pytest tests/tools/test_export_onnx.py` 验证 ORT vs PyTorch 数值一致 **阶段 4 — 库代码改写(纯 numpy + ORT)** - 抽 `_postprocess.py`、复制 `_coord.py`、写 `_assets.py`、`mouse.py`、`scroll.py`、`errors.py` - 重写 `__init__.py` 暴露 API + lru_cache - 删除阶段 1~2 残留的 torch 引用源文件 - 验收: golden 测试通过;`uv pip install .` 不引入 torch;`from ai_mouse import generate` 直接出结果 **阶段 5 — 文档与清理** - 重写 README(SDK 视角) - 新建 CHANGELOG (0.2.0) - 更新 CLAUDE.md - 删 `JointCVAE` - 写 `examples/quickstart.py` #### 风险与对策 | 风险 | 对策 | |---|---| | 权重二进制进 git (~3 MB) | 直接 commit;不引入 LFS | | 中间状态 CI 不绿 | 每个阶段独立 PR,自包含通过 | | `tools/` 内部 import 风格 | 一律绝对 `from tools.X`;`tools/__init__.py` 留空 | | 阶段 4 改 API 的"半破"中间态 | 阶段 4 是唯一签名变化节点,同 PR 内同步改 README/CLAUDE/examples | | `data/models_v2/` 路径引用混乱 | 库代码不读 `data/`(只读 `importlib.resources`);仅 `tools/export_onnx.py` 和 `tools/trainer.py` 关心 `data/` | | Golden 测试容差选择 | 起始 `atol=2`(像素 + ms),实测后可收紧 | | Windows 路径 | 全程 `pathlib.Path`;`importlib.resources` 跨平台 | | ORT CPU/GPU 包互斥 | README 明示 `pip install onnxruntime-gpu` 替换默认 `onnxruntime` | ### §5 测试策略 #### 两层测试,依赖边界硬隔离 `tests/unit/` 只装 `numpy + onnxruntime` 必须通过 — "纯推理 SDK"的硬保证。 `tests/tools/` 装全套 `[dev]` 才跑训练/服务器/eval 测试。 #### Golden 测试(锁住语义不漂移) 迁移**开始之前**用旧 torch 实现跑 8 个 case × 4 个 seed = 32 个固定输出,存 `tests/unit/data/golden_mouse.npz`。 ```python CASES = [ ((100, 200), (900, 400)), # 水平 800px ((500, 500), (500, 100)), # 垂直 400px ((200, 600), (800, 200)), # 720px 对角 ((100, 100), (130, 110)), # 极短 31px ((50, 50), (1500, 900)), # 极长 1700px ((400, 300), (500, 300)), # 水平 100px ((300, 300), (700, 700)), # 45° 对角 ((600, 400), (200, 100)), # 反向对角 ] SEEDS = (0, 1, 2, 3) ``` 迁移后断言 `np.allclose(new_pts, golden_pts, atol=2)`(像素 + ms 容差)。 Scroll 同样做 golden;容差以 deltaY 量化粒度允许 ±1 个 quantum。 #### ONNX vs PyTorch 数值一致性 `tests/tools/test_export_onnx.py`:训练 1-epoch toy 模型(seq_len=8, d_model=16)⇒ 导出 ⇒ 同输入跑 PyTorch & ORT ⇒ `assert np.allclose(atol=1e-4)`。 真模型的一致性由 `tools/export_onnx.py` 末尾烟雾测试在导出时执行。 #### `_postprocess.py` 单元测试 每个纯函数单独测,边界条件覆盖。详细测试矩阵见原讨论;关键覆盖点: | 函数 | 关键断言 | |---|---| | `snap_endpoints` | 输出端点严格 = (0,0) 和 (1,0);前 (seq_len - n_snap) 个点不变 | | `enforce_forward_monotonic` | 输出非递减 | | `gaussian_smooth` | 端点保持;权重和归一 | | `build_timestamps` | sum ≈ total_duration;严格单调;首项 0 | | `sample_duration` | 1000 次采样中位数落在 bin 对应 `[0.5×, 2×]exp(mu_log)` | | `truncnorm_sample` | 落在 `[low, high]`;均值 ≈ mu;兜底分支独立覆盖 | | `resample_arc` | 长度变换正确;N=M 恒等 | #### `MouseModel` / `ScrollModel` 行为测试 - session 复用(同实例多次 generate 不重建 session) - 顶层 `generate()` 单例缓存命中 - `seed` 复现性(同 seed 两次调用结果一致) - `click=False` 返回纯轨迹 - 非法 `model_path` 抛 `ModelLoadError` #### 线程安全测试 `ThreadPoolExecutor` 并发 32 次 `m.generate(...)`,断言全部完成、无异常。不保证并发下同 seed 严格一致(ORT 内部调度)。 #### CI 矩阵 ```yaml jobs: library: # 只装运行时依赖 + pytest — 拦截库代码偷偷 import torch runs: | uv venv && uv pip install -e . pytest uv run pytest tests/unit -v dev: runs: | uv sync --group dev uv run pytest tests/ -v ``` 注: `pytest` 不能进运行时依赖(下游用户不需要),只能在 CI 安装行里临时加。 `pytest-asyncio` 仅 `tests/tools/test_server.py` 需要,留在 `[dev]` 群组。 OS 矩阵: Linux + Windows;Python 3.12 + 3.13。 #### 不做的事 - ❌ 不测视觉真实性(主观,留给人眼审 eval report) - ❌ 不跑真 web E2E(httpx ASGI 集成已足够) - ❌ CI 不重训模型(用已 commit 的 ONNX) ### §6 文档与下游集成 #### README.md 重写(SDK 视角) 读者画像从"训练者"切换到"集成者"。结构: 1. 标题与一句话定位 2. 安装(`pip install git+...` + GPU 可选说明) 3. 快速上手(鼠标、滚轮、类用法、provider 切换、复现性) 4. API 概览(表格) 5. 线程安全说明 6. 自训练模型 — 跳转 Development 小节 7. Development(训练命令、Web UI、ONNX 导出 — 目标读者是 contributor) #### `examples/quickstart.py` 最小可跑示例:`from ai_mouse import generate` ⇒ 打印 6 行结果 + 一段"实际使用模式"伪代码(时间戳节流 + 占位的鼠标移动调用)。 #### CHANGELOG.md(新建,Keep a Changelog 格式) 锚定 0.2.0: ``` ## [0.2.0] - YYYY-MM-DD ### Changed (breaking) - Inference 不再依赖 PyTorch;运行时仅 numpy + onnxruntime - 公开 API 新增 MouseModel / ScrollModel 类 - 函数式 generate / generate_scroll 签名: keyword-only;新增 click=, seed=;移除 config= - 训练 / 服务器 / 采集 / 评估代码移至 tools/,不再打入 wheel ### Added - Wheel 内置 ONNX 预训练权重 (~3 MB) - ORT provider 自定义 (GPU/DirectML) ### Removed - JointCVAE (legacy) - ai_mouse.config.GenerateConfig 顶层导出 ``` #### CLAUDE.md 更新清单 - Commands: `python -m ai_mouse ` ⇒ `python -m tools ` - 加一节"Library vs tools 边界": 库代码禁止 `import torch` - 移除 "Bundled weights live in `data/models_v2/`" ⇒ 改 `src/ai_mouse/assets/` - 测试章节区分 `tests/unit` / `tests/tools` #### 下游契约 — 成功验收标准 1. `pip install git+...` ⇒ `from ai_mouse import generate` 直接出结果,无缺失文件、无 GB 级依赖 2. `import ai_mouse` 时间 < 200 ms 3. 单次 `generate()` < 50 ms (CPU,Euler 10 步 + 后处理) #### 不做的事 - ❌ 不写 API reference 网站(docstring 足够) - ❌ 不写迁移 guide(破坏面有限,CHANGELOG 一段话) - ❌ 不录视频 / GIF - ❌ 不引入 docs i18n 流程 ## 开放问题 无。所有决策点已经在与用户的对话中明确。 ## 后续行动 本 spec 经用户批准后,进入 implementation plan 阶段(`writing-plans` skill),把 §4 的 5 个迁移阶段展开为可执行的 task list。