chore: initialize git repo, add matplotlib dep, extend config
- Add .gitignore for Python/data/models - Add matplotlib>=3.8.0 for eval plots - Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses
This commit is contained in:
3151
docs/superpowers/plans/2026-05-10-balabit-pretrain-refactor.md
Normal file
3151
docs/superpowers/plans/2026-05-10-balabit-pretrain-refactor.md
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,393 @@
|
||||
# Balabit 预训练 + Fine-tune 重构设计
|
||||
|
||||
**日期**：2026-05-10
|
||||
**作者**：Claude + 用户
|
||||
**状态**：待用户评审
|
||||
|
||||
## 1. 背景与目标
|
||||
|
||||
### 1.1 问题
|
||||
当前 `TrajectoryFlowModel` 仅用 605 条用户采集数据训练，生成质量差：
|
||||
- **空间锯齿**：lateral 方向高频抖动（见 verify 页面截图）
|
||||
- **时间模板化**：5 条生成轨迹的 Δt 曲线几乎完全重合（多样性丧失）
|
||||
- 反检测能力弱
|
||||
|
||||
### 1.2 根本原因
|
||||
1. **数据量不足**（605 条），模型欠拟合 + 记噪声
|
||||
2. **后处理过度**：[generator.py](../../../ai_mouse/generator.py) 中的 `speed_profile`（确定性钟形曲线）和 `median±1.1` 硬 clip 把多样性压扁
|
||||
|
||||
### 1.3 目标
|
||||
- 用 Balabit Mouse Dynamics Challenge 数据集（公开真实用户会话）做预训练
|
||||
- 用现有 605 条做 fine-tune 适配本任务分布
|
||||
- 重写后处理，让真人数据本身的速度模式自然显现
|
||||
- 输出量化评估报表，便于迭代判断进步
|
||||
|
||||
### 1.4 不在范围（YAGNI）
|
||||
- ❌ 滚轮模型 `ai_mouse/scroll/`（独立子系统，本次冻结不动）
|
||||
- ❌ 模型架构变更（仍是 `TrajectoryFlowModel`，不换 Diffusion——CPU 推理约束）
|
||||
- ❌ 对抗训练 / Discriminator（验证完基础方案再说）
|
||||
- ❌ 浏览器插件、其他公开数据集
|
||||
- ❌ 前端 UI 变更
|
||||
|
||||
## 2. 整体架构
|
||||
|
||||
```
|
||||
Balabit raw sessions (CSV)
|
||||
│
|
||||
▼ [新] ai_mouse/data_adapters/balabit.py
|
||||
click 锚定切分（Pressed 事件前 W ms 的 Move）
|
||||
│
|
||||
▼
|
||||
data/pretrain_traces.jsonl ← 与现有 traces.jsonl 格式 100% 兼容
|
||||
│
|
||||
▼ [改] ai_mouse/trainer.py（streaming dataloader、resume_from）
|
||||
data/models_v2_pretrained/ ← 预训练 checkpoint（Balabit only）
|
||||
│
|
||||
▼ [改] trainer.py 加 --resume-from 支持加载已有权重
|
||||
data/models_v2/ ← 605 条 fine-tune 后的最终权重（部署用）
|
||||
│
|
||||
▼ [改] ai_mouse/generator.py（砍硬模板后处理）
|
||||
推理：CPU < 200ms
|
||||
|
||||
评估：[新] ai_mouse/eval/
|
||||
▼
|
||||
data/eval_reports/YYYY-MM-DD-<tag>.md ← 量化指标 + 直方图 + FFT
|
||||
```
|
||||
|
||||
**关键设计原则**：
|
||||
- **数据格式不变**：Balabit 转换后产出和现有 `traces.jsonl` 格式 100% 一致。`load_and_prepare_data` / 旋转坐标系 / 6× 增强逻辑全部不动
|
||||
- **训练流程加一个阶段**：原来"605 → train → models_v2"，现在变成"Balabit → pretrain → models_v2_pretrained → 605 fine-tune → models_v2"
|
||||
- **滚轮子系统完全不动**
|
||||
|
||||
## 3. 文件变更清单
|
||||
|
||||
### 3.1 新增
|
||||
|
||||
| 路径 | 用途 |
|
||||
|---|---|
|
||||
| `ai_mouse/data_adapters/__init__.py` | 包初始化 |
|
||||
| `ai_mouse/data_adapters/balabit.py` | Balabit CSV → traces.jsonl 适配器 + CLI |
|
||||
| `ai_mouse/eval/__init__.py` | 包初始化 |
|
||||
| `ai_mouse/eval/metrics.py` | 运动学指标计算（速度/加速度/jerk/FFT） |
|
||||
| `ai_mouse/eval/report.py` | Markdown 报表生成 |
|
||||
| `ai_mouse/eval/__main__.py` | CLI: `python -m ai_mouse.eval` |
|
||||
| `ai_mouse/__main__.py` | CLI: `python -m ai_mouse train ...`（统一入口） |
|
||||
| `tests/test_balabit_adapter.py` | 适配器单元测试 |
|
||||
| `tests/test_eval_metrics.py` | 指标计算正确性测试 |
|
||||
|
||||
### 3.2 修改
|
||||
|
||||
| 路径 | 改动 |
|
||||
|---|---|
|
||||
| `ai_mouse/trainer.py` | 加 `resume_from` 参数；增强从一次性 numpy 改成 Dataset 内 on-the-fly；可选 AMP（有 GPU 时） |
|
||||
| `ai_mouse/generator.py` | 砍 `speed_profile` 与 `median±1.1` clip；新增 5 点高斯 lateral 平滑 |
|
||||
| `ai_mouse/config.py` | 新增 `PretrainConfig`、`FinetuneConfig`、`BalabitAdapterConfig`、`EvalConfig` |
|
||||
| `tests/test_generator.py` | 适配新后处理 |
|
||||
| `tests/test_trainer.py` | 加 resume_from 测试 |
|
||||
|
||||
### 3.3 冻结不动
|
||||
|
||||
`ai_mouse/scroll/*`、`ai_mouse/coord.py`、`ai_mouse/models.py`、`ai_mouse/utils.py`、`ai_mouse/collector.py`、`ai_mouse/server/*`、`static/*`、`tests/test_scroll_*.py`、`tests/test_coord.py`、`tests/test_models.py`、`tests/test_server.py`
|
||||
|
||||
## 4. Balabit 适配器（`ai_mouse/data_adapters/balabit.py`）
|
||||
|
||||
### 4.1 输入格式
|
||||
Balabit Mouse Dynamics Challenge 每个 session 是 CSV 文件，列：
|
||||
```
|
||||
record timestamp, client timestamp, button, state, x, y
|
||||
```
|
||||
- `state` ∈ {`Move`, `Pressed`, `Released`, `Drag`, `Scroll`}
|
||||
- `button` ∈ {`NoButton`, `Left`, `Right`}
|
||||
|
||||
### 4.2 切分逻辑（click 锚定）
|
||||
1. 扫描 session 所有事件
|
||||
2. 找每个 `Pressed` 事件 P
|
||||
3. 回溯前 W ms（默认 `W=1200`）内的所有 `Move` 事件，构成一段 trace
|
||||
4. 段的 `start = 第一个 Move 的 (x,y)`，`end = P 的 (x,y)`
|
||||
5. 时间戳归零（第一个 Move 的 t=0）
|
||||
|
||||
### 4.3 过滤规则
|
||||
**丢弃整段**（不修复，直接丢弃这一条 trace），满足任一条件即丢：
|
||||
- Move 事件数 < 5
|
||||
- `dist(start, end) < 50` px
|
||||
- 时间跨度 > 5000ms（避免长停顿）
|
||||
- `start` 或 `end` 任一坐标 < 0 或 > 5000（避免跨屏瞬移异常值）
|
||||
- 总弧长 < 50 px（避免抖动残留）
|
||||
- 段内任意相邻 Move 之间时间差 > 200ms（采样断档，宁可整段丢也不要含断点的样本）
|
||||
|
||||
### 4.4 输出
|
||||
追加到 `data/pretrain_traces.jsonl`，每行：
|
||||
```json
|
||||
{"meta":{"start":[x,y],"end":[x,y],"dist":int,"angle":float,"source":"balabit","session_id":"user7_session_42"},"events":[{"type":"move","x":int,"y":int,"t":int}, ...]}
|
||||
```
|
||||
|
||||
**关于 click events 的兼容性**：Balabit 转换后的 trace **不包含** `down`/`up` 事件（预训练只学移动）。这与现有 `trainer.py:139-142` 的逻辑兼容——`load_and_prepare_data` 用 `if downs and ups` 检查，没有就跳过。后果：
|
||||
- 预训练阶段产出的 `click_dist.json` 会基于零样本，**写入默认值**（mu=80, sigma=30）
|
||||
- Fine-tune 阶段重新基于 605 条产出真正的 `click_dist.json`，覆盖默认值
|
||||
- `flow_model.pt` 和 `train_config.json` 是预训练真正要保留的产物
|
||||
|
||||
### 4.5 CLI
|
||||
```bash
|
||||
uv run python -m ai_mouse.data_adapters.balabit \
|
||||
--input /path/to/balabit/sessions/ \
|
||||
--output data/pretrain_traces.jsonl \
|
||||
--window-ms 1200
|
||||
```
|
||||
|
||||
## 5. 训练 pipeline 改造（`ai_mouse/trainer.py`）
|
||||
|
||||
### 5.1 新增参数
|
||||
```python
|
||||
def train(
|
||||
data_path: Path,
|
||||
output_dir: Path,
|
||||
epochs: int = 300,
|
||||
batch_size: int = 64,
|
||||
lr: float = 3e-4,
|
||||
seq_len: int = 64,
|
||||
progress_callback: Callable[[dict], None] | None = None,
|
||||
config: TrainConfig | None = None,
|
||||
resume_from: Path | None = None, # 新增：加载已有权重
|
||||
use_amp: bool = False, # 新增：mixed precision
|
||||
) -> None:
|
||||
```
|
||||
|
||||
### 5.2 Dataset 改造
|
||||
当前 `_augment` 是**一次性把全部数据 6× 复制到内存**。Balabit 几万条样本 × 6 后内存占用大。改为：
|
||||
|
||||
```python
|
||||
class TrajectoryDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, seq, cond, augment: bool = True):
|
||||
self.seq = seq
|
||||
self.cond = cond
|
||||
self.augment = augment
|
||||
self._n_aug = 6 if augment else 1
|
||||
|
||||
def __len__(self):
|
||||
return len(self.seq) * self._n_aug
|
||||
|
||||
def __getitem__(self, idx):
|
||||
base_idx = idx // self._n_aug
|
||||
aug_id = idx % self._n_aug
|
||||
seq, cond = self.seq[base_idx], self.cond[base_idx]
|
||||
return self._apply_augment(seq, cond, aug_id)
|
||||
```
|
||||
|
||||
每次 `__getitem__` 即时增强，内存占用 = 原始数据 1×。
|
||||
|
||||
### 5.3 Resume from checkpoint
|
||||
```python
|
||||
if resume_from is not None:
|
||||
state = torch.load(resume_from / "flow_model.pt", weights_only=True)
|
||||
model.load_state_dict(state)
|
||||
```
|
||||
|
||||
如果 fine-tune 阶段的 `cond_dim` 与预训练不一致，需要明确报错。
|
||||
|
||||
### 5.4 完整两阶段训练流程
|
||||
```bash
|
||||
# 阶段 1：预训练
|
||||
uv run python -m ai_mouse train \
|
||||
--data data/pretrain_traces.jsonl \
|
||||
--output data/models_v2_pretrained \
|
||||
--epochs 200 --lr 3e-4 --batch-size 128
|
||||
|
||||
# 阶段 2：微调
|
||||
uv run python -m ai_mouse train \
|
||||
--data data/traces.jsonl \
|
||||
--output data/models_v2 \
|
||||
--epochs 50 --lr 1e-5 --batch-size 64 \
|
||||
--resume-from data/models_v2_pretrained
|
||||
```
|
||||
|
||||
### 5.5 服务端 API
|
||||
`POST /api/train` 当前不支持 resume。**本次不改服务端 API**——CLI 是主要预训练入口，UI 上的"训练"按钮仍用于 605 条的 fine-tune（默认会自动从 `models_v2_pretrained` resume，如果存在）。
|
||||
|
||||
逻辑：
|
||||
- `models_v2_pretrained/flow_model.pt` 存在 → fine-tune 模式（lr=1e-5, epochs=50）
|
||||
- 不存在 → 走原逻辑（from scratch）
|
||||
|
||||
## 6. Generator 改造（`ai_mouse/generator.py`）
|
||||
|
||||
### 6.1 砍掉
|
||||
- **lines 263–271** `max_allowed = median + 1.1`、`min_allowed = median - 1.1` 硬 clip 整段
|
||||
- **lines 273–286** `speed_profile` 整段（acceleration/deceleration phase）
|
||||
|
||||
### 6.2 保留
|
||||
- 端点 snap (lines 220–229)
|
||||
- 起点 lateral 衰减 (lines 232–236)
|
||||
- forward 单调性强制 (lines 239–246)
|
||||
- log_dt 安全 clip [0, 5] (line 255)
|
||||
- click duration 采样
|
||||
|
||||
### 6.3 新增：lateral 5 点高斯平滑
|
||||
```python
|
||||
# 在 lateral monotonic fix 之后、decode_trajectory 之前
|
||||
def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
|
||||
"""5-point gaussian smoothing, preserving endpoints."""
|
||||
kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
|
||||
kernel /= kernel.sum()
|
||||
smoothed = np.convolve(x, kernel, mode="same")
|
||||
smoothed[0] = x[0] # preserve endpoint
|
||||
smoothed[-1] = x[-1]
|
||||
return smoothed
|
||||
|
||||
lateral = _gaussian_smooth(lateral, sigma=1.0)
|
||||
```
|
||||
|
||||
只对 lateral 平滑，**不对 forward 平滑**（forward 已经被单调性约束保护，再平滑会损害准确性）。
|
||||
|
||||
### 6.4 预期效果
|
||||
- 解决高频锯齿（lateral 平滑）
|
||||
- Δt 多样性恢复（不再被压扁到 median ± 1.1）
|
||||
- 速度模式由模型自己决定（learned from Balabit real data）
|
||||
|
||||
## 7. 评估模块（`ai_mouse/eval/`）
|
||||
|
||||
### 7.1 指标（`metrics.py`）
|
||||
|
||||
每条轨迹计算：
|
||||
- **速度** `v[i] = sqrt(dx² + dy²) / dt`（单位 px/ms）
|
||||
- **加速度** `a[i] = (v[i+1] - v[i]) / dt`
|
||||
- **Jerk** `j[i] = (a[i+1] - a[i]) / dt`
|
||||
- **Δt 序列**
|
||||
|
||||
聚合到样本集：
|
||||
- 速度/加速度/jerk 的均值、std、变异系数 CV、p25/p50/p75/p95
|
||||
- Δt 分布的 CV
|
||||
- **FFT 频谱**：每条轨迹 lateral 信号做 FFT，看 4–12Hz 频段是否有 peak（生理震颤）
|
||||
- **多样性**：N 条样本之间的 PCA 方差贡献（衡量是否模板化）
|
||||
|
||||
对比"参考分布"（从 `pretrain_traces.jsonl` 随机抽 1000 条作为 holdout）：
|
||||
- 速度/加速度分布的 KL 散度（直方图离散化估计）
|
||||
- jerk 分布的 KL 散度
|
||||
|
||||
### 7.2 报表（`report.py`）
|
||||
|
||||
输出 Markdown 到 `data/eval_reports/YYYY-MM-DD-<tag>.md`，结构：
|
||||
|
||||
```markdown
|
||||
# Eval Report: <tag> (2026-05-10 18:30)
|
||||
|
||||
## 模型信息
|
||||
- Checkpoint: data/models_v2/flow_model.pt
|
||||
- 训练参数: ...
|
||||
|
||||
## 摘要
|
||||
| 指标 | 生成 | 参考 | 评价 |
|
||||
|---|---|---|---|
|
||||
| 速度 KL | 0.12 | 0 | OK |
|
||||
| FFT 4-12Hz peak | 8.3Hz @ 0.04 | 7.1Hz @ 0.05 | OK |
|
||||
| Δt CV | 0.45 | 0.52 | 接近 |
|
||||
| ...
|
||||
|
||||
## 直方图（PNG 嵌入）
|
||||
|
||||
...
|
||||
|
||||
## 5 条生成轨迹示例
|
||||
|
||||
|
||||
## vs 上次报表
|
||||
- 速度 KL 从 0.34 → 0.12（提升）
|
||||
- ...
|
||||
```
|
||||
|
||||
PNG 用 matplotlib 输出到 `data/eval_reports/plots/`。
|
||||
|
||||
### 7.3 CLI
|
||||
```bash
|
||||
uv run python -m ai_mouse.eval \
|
||||
--model-dir data/models_v2 \
|
||||
--reference data/pretrain_traces.jsonl \
|
||||
--n-samples 1000 \
|
||||
--output data/eval_reports/2026-05-10-baseline.md \
|
||||
--tag baseline
|
||||
```
|
||||
|
||||
## 8. 配置变更（`ai_mouse/config.py`）
|
||||
|
||||
新增：
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class PretrainConfig:
|
||||
"""Hyperparameters for Balabit pretraining."""
|
||||
epochs: int = 200
|
||||
batch_size: int = 128
|
||||
lr: float = 3e-4
|
||||
seq_len: int = 64
|
||||
|
||||
@dataclass
|
||||
class FinetuneConfig:
|
||||
"""Hyperparameters for fine-tuning on user data."""
|
||||
epochs: int = 50
|
||||
batch_size: int = 64
|
||||
lr: float = 1e-5 # 比预训练小一个数量级
|
||||
seq_len: int = 64
|
||||
|
||||
@dataclass
|
||||
class BalabitAdapterConfig:
|
||||
"""Settings for Balabit data conversion."""
|
||||
window_ms: int = 1200
|
||||
min_dist: int = 50
|
||||
min_events: int = 5
|
||||
max_span_ms: int = 5000
|
||||
max_gap_ms: int = 200
|
||||
|
||||
@dataclass
|
||||
class EvalConfig:
|
||||
"""Settings for evaluation report generation."""
|
||||
n_samples: int = 1000
|
||||
fft_freq_band: tuple[float, float] = (4.0, 12.0)
|
||||
kl_bins: int = 50
|
||||
```
|
||||
|
||||
`TrainConfig` 保持不变（向后兼容现有训练脚本）。
|
||||
|
||||
## 9. 测试策略
|
||||
|
||||
### 9.1 新增测试
|
||||
- `test_balabit_adapter.py`：用合成 CSV 测试切分逻辑、过滤规则、边界条件（空 session、无 click、坐标溢出）
|
||||
- `test_eval_metrics.py`：固定输入下指标计算的正确性
|
||||
|
||||
### 9.2 更新测试
|
||||
- `test_generator.py`：移除对 `speed_profile` 的断言；新增 lateral 平滑断言
|
||||
- `test_trainer.py`：新增 `resume_from` 测试
|
||||
|
||||
### 9.3 保持通过
|
||||
所有非 `test_generator.py`、非 `test_trainer.py` 的测试保持通过。`test_server.py` 不变（服务端 API 未改）。
|
||||
|
||||
## 10. 验收标准
|
||||
|
||||
最终评估报表（`data/eval_reports/<final>.md`）应显示：
|
||||
|
||||
1. **主观**：5 条生成轨迹的 Δt 曲线明显多样化（不再重合）
|
||||
2. **主观**：lateral 无高频锯齿
|
||||
3. **量化**：速度分布 KL（vs Balabit holdout）< 当前实现（pre-refactor baseline，第一次跑评估时记下）的 50%
|
||||
4. **量化**：FFT 频谱在 4–12Hz 区间出现 peak
|
||||
5. **回归**：所有非废弃测试保持通过
|
||||
|
||||
## 11. 工作量估计
|
||||
|
||||
| 阶段 | 时间 |
|
||||
|---|---|
|
||||
| Balabit 适配器（含测试） | 1–2 天 |
|
||||
| Trainer 改造（streaming + resume） | 1 天 |
|
||||
| Generator 后处理改造 | 0.5 天 |
|
||||
| 评估模块 | 1–2 天 |
|
||||
| 跑预训练 + fine-tune（CPU/GPU 视情况） | 0.5–2 天 |
|
||||
| 调参迭代 + 报表对比 | 1–3 天 |
|
||||
|
||||
**最小可行版本**（搭起来跑通第一版）：3–5 天
|
||||
**完整调到验收**：1–2 周
|
||||
|
||||
## 12. 风险与备选
|
||||
|
||||
| 风险 | 缓解 |
|
||||
|---|---|
|
||||
| Balabit 切分后样本数不足（< 5000） | 放宽 `min_dist` 到 30，或扩大窗口到 2000ms |
|
||||
| 预训练后 fine-tune 出现灾难性遗忘 | lr 调更小（1e-6），epochs 减到 20 |
|
||||
| 模型架构 cond_dim 在预训练/fine-tune 不一致 | 强制相同；不一致时直接 raise |
|
||||
| 评估报表实现工作量过大 | 第一版只做基础指标（速度/Δt CV），FFT 和 PCA 后置 |
|
||||
| Balabit 数据集合规问题 | 仅本地使用，不分发，不商用 |
|
||||
Reference in New Issue
Block a user