Huang Qi
4d414fd180
- Add .gitignore for Python/data/models - Add matplotlib>=3.8.0 for eval plots - Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses
394 lines
14 KiB
Markdown
394 lines
14 KiB
Markdown
# 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 数据集合规问题 | 仅本地使用，不分发，不商用 |
|