Huang Qi
4d414fd180
- Add .gitignore for Python/data/models - Add matplotlib>=3.8.0 for eval plots - Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses
161 lines
4.4 KiB
Python
161 lines
4.4 KiB
Python
"""Centralised configuration for ai_mouse.
|
|
|
|
All magic numbers and hyperparameters live here so they can be tuned
|
|
from one place, overridden per-instance, or serialised to JSON.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Training configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class TrainConfig:
|
|
"""Hyperparameters for Flow Matching training."""
|
|
|
|
epochs: int = 300
|
|
batch_size: int = 64
|
|
lr: float = 3e-4
|
|
seq_len: int = 64
|
|
# Transformer backbone
|
|
d_model: int = 128
|
|
nhead: int = 4
|
|
num_layers: int = 4
|
|
dim_feedforward: int = 256
|
|
dropout: float = 0.1
|
|
# Flow matching
|
|
sigma_min: float = 1e-4
|
|
# Data augmentation
|
|
augment: bool = True
|
|
# Duration distribution bins
|
|
dist_bins: list[float] = field(
|
|
default_factory=lambda: [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")]
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Generation configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class GenerateConfig:
|
|
"""Tuneable knobs for Flow Matching inference."""
|
|
|
|
n_steps: int = 10 # Euler ODE steps
|
|
seq_len: int = 64
|
|
dt_clip_min_ms: float = 2.0
|
|
dt_clip_max_ms: float = 150.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scroll subsystem configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class ScrollModeConfig:
|
|
"""Parameters for a single scroll collection mode."""
|
|
|
|
dist_min: int
|
|
dist_max: int
|
|
success_radius: int
|
|
|
|
|
|
SCROLL_MODES: dict[str, ScrollModeConfig] = {
|
|
"target": ScrollModeConfig(dist_min=500, dist_max=3000, success_radius=80),
|
|
"fast": ScrollModeConfig(dist_min=3000, dist_max=8000, success_radius=120),
|
|
"precise": ScrollModeConfig(dist_min=200, dist_max=800, success_radius=40),
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ScrollTrainConfig:
|
|
"""Hyperparameters for ScrollCVAE training."""
|
|
|
|
epochs: int = 100
|
|
batch_size: int = 32
|
|
lr: float = 5e-4
|
|
seq_len: int = 32
|
|
beta_max: float = 0.3
|
|
beta_warmup_epochs: int = 15
|
|
weight_delta: float = 1.0
|
|
weight_log_dt: float = 1.5
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Server configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class ServerConfig:
|
|
"""Web server and data directory settings."""
|
|
|
|
host: str = "127.0.0.1"
|
|
port: int = 8765
|
|
data_dir: Path = field(default_factory=lambda: Path("data"))
|
|
canvas_width: int = 800
|
|
canvas_height: int = 600
|
|
open_browser: bool = True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pretraining (Balabit) configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class PretrainConfig:
|
|
"""Hyperparameters for Balabit pretraining stage."""
|
|
|
|
epochs: int = 200
|
|
batch_size: int = 128
|
|
lr: float = 3e-4
|
|
seq_len: int = 64
|
|
|
|
|
|
@dataclass
|
|
class FinetuneConfig:
|
|
"""Hyperparameters for fine-tuning on user-collected data."""
|
|
|
|
epochs: int = 50
|
|
batch_size: int = 64
|
|
lr: float = 1e-5 # 比预训练小一个数量级，防止灾难性遗忘
|
|
seq_len: int = 64
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Balabit adapter configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class BalabitAdapterConfig:
|
|
"""Settings for Balabit CSV → traces.jsonl conversion."""
|
|
|
|
window_ms: int = 1200 # click 前回溯窗口
|
|
min_dist: int = 50 # 最小起终点距离 (px)
|
|
min_events: int = 5 # 最小 Move 事件数
|
|
max_span_ms: int = 5000 # 最大段时间跨度 (ms)
|
|
max_gap_ms: int = 200 # 段内相邻 Move 最大时间差
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Eval configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class EvalConfig:
|
|
"""Settings for evaluation report generation."""
|
|
|
|
n_samples: int = 1000
|
|
fft_freq_band: tuple[float, float] = (4.0, 12.0) # 生理震颤频段 (Hz)
|
|
kl_bins: int = 50
|