refactor: move trainer/models/utils/config to tools/
This commit is contained in:
509
tools/trainer.py
Normal file
509
tools/trainer.py
Normal file
@@ -0,0 +1,509 @@
|
||||
"""Training pipeline for Conditional Flow Matching mouse trajectory model.
|
||||
|
||||
Pipeline:
|
||||
1. Load traces from JSONL, convert to rotated coordinate frame
|
||||
2. Apply 6× data augmentation
|
||||
3. Train TrajectoryFlowModel with OT-Conditional Flow Matching:
|
||||
- x1 = real data, x0 = randn_like(x1), t = rand(B)
|
||||
- x_t = (1-t)*x0 + t*x1
|
||||
- v_target = x1 - x0
|
||||
- v_pred = model(x_t, t, cond)
|
||||
- loss = MSE(v_pred, v_target)
|
||||
4. Save: flow_model.pt, click_dist.json, duration_dist.json, train_config.json
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from ai_mouse.coord import encode_trajectory
|
||||
from tools.config import TrainConfig
|
||||
from tools.models import TrajectoryFlowModel
|
||||
from tools.utils import resample_arc
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Distance bins for duration distribution (in pixels)
|
||||
_DIST_BINS: list[float] = [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data loading
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_and_prepare_data(
|
||||
data_path: Path,
|
||||
seq_len: int = 64,
|
||||
) -> tuple[np.ndarray, np.ndarray, list[float]]:
|
||||
"""Load JSONL traces and convert to rotated-frame tensors.
|
||||
|
||||
Args:
|
||||
data_path: path to traces.jsonl
|
||||
seq_len: number of time steps to resample each trajectory to
|
||||
|
||||
Returns:
|
||||
seq: (N, seq_len, 3) float32 — (forward, lateral, log_Δt)
|
||||
cond: (N, 3) float32 — [dist_norm, log_dist, log_dur]
|
||||
click_durs: list of float click durations in ms
|
||||
"""
|
||||
data_path = Path(data_path)
|
||||
seq_list: list[np.ndarray] = []
|
||||
cond_list: list[np.ndarray] = []
|
||||
click_durs: list[float] = []
|
||||
|
||||
for i, raw_line in enumerate(data_path.read_text(encoding="utf-8").splitlines(), 1):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
trace = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
meta = trace["meta"]
|
||||
events = trace["events"]
|
||||
|
||||
if "start" not in meta or "end" not in meta:
|
||||
logger.warning("Skipping line %d: missing start/end in meta", i)
|
||||
continue
|
||||
|
||||
sx, sy = meta["start"]
|
||||
ex, ey = meta["end"]
|
||||
|
||||
# Extract move events
|
||||
moves = [(e["x"], e["y"], e["t"]) for e in events if e["type"] == "move"]
|
||||
if len(moves) < 2:
|
||||
continue
|
||||
|
||||
xs = np.array([m[0] for m in moves], dtype=float)
|
||||
ys = np.array([m[1] for m in moves], dtype=float)
|
||||
ts = np.array([m[2] for m in moves], dtype=float)
|
||||
|
||||
xy_raw = np.stack([xs, ys], axis=1)
|
||||
|
||||
# Reject degenerate (zero-length) trajectories
|
||||
total_arc = float(np.linalg.norm(np.diff(xy_raw, axis=0), axis=1).sum())
|
||||
if total_arc < 1.0:
|
||||
continue
|
||||
|
||||
# Resample spatial positions to seq_len via arc-length
|
||||
xy_resampled = resample_arc(xy_raw, seq_len) # (seq_len, 2)
|
||||
|
||||
# Resample timestamps along the same arc-length grid
|
||||
arc_dist = np.concatenate(
|
||||
[[0.0], np.cumsum(np.linalg.norm(np.diff(xy_raw, axis=0), axis=1))]
|
||||
)
|
||||
s_uniform = np.linspace(0.0, arc_dist[-1], seq_len)
|
||||
ts_resampled = np.interp(s_uniform, arc_dist, ts)
|
||||
|
||||
# Convert spatial coords to rotated frame (forward, lateral)
|
||||
fl = encode_trajectory(xy_resampled, (sx, sy), (ex, ey)) # (seq_len, 2)
|
||||
|
||||
# Compute Δt intervals (length seq_len-1) → log(Δt+1), pad 0 at front
|
||||
dt_raw = np.diff(ts_resampled).clip(0.0)
|
||||
log_dt = np.log(dt_raw + 1.0) # (seq_len-1,)
|
||||
log_dt_padded = np.concatenate([[0.0], log_dt]) # (seq_len,) — first step has no interval
|
||||
|
||||
# Stack into (seq_len, 3)
|
||||
seq_arr = np.stack([fl[:, 0], fl[:, 1], log_dt_padded], axis=1).astype(np.float32)
|
||||
|
||||
# Condition vector
|
||||
dist = float(meta["dist"]) if meta["dist"] > 0 else float(
|
||||
math.hypot(ex - sx, ey - sy)
|
||||
)
|
||||
dist = max(dist, 1.0)
|
||||
total_dur = float(ts_resampled[-1] - ts_resampled[0])
|
||||
total_dur = max(total_dur, 1.0)
|
||||
|
||||
cond_arr = np.array(
|
||||
[
|
||||
dist / 2000.0, # dist_norm
|
||||
math.log(dist / 100.0), # log_dist
|
||||
math.log(total_dur / 500.0), # log_dur
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
seq_list.append(seq_arr)
|
||||
cond_list.append(cond_arr)
|
||||
|
||||
# Click duration (down→up)
|
||||
downs = [e for e in events if e["type"] == "down"]
|
||||
ups = [e for e in events if e["type"] == "up"]
|
||||
if downs and ups:
|
||||
click_durs.append(float(ups[-1]["t"] - downs[-1]["t"]))
|
||||
|
||||
if not seq_list:
|
||||
raise ValueError(f"No valid traces found in {data_path}")
|
||||
|
||||
return (
|
||||
np.stack(seq_list, axis=0), # (N, seq_len, 3)
|
||||
np.stack(cond_list, axis=0), # (N, 3)
|
||||
click_durs,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data augmentation (6×)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _augment(
|
||||
seq: np.ndarray,
|
||||
cond: np.ndarray,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""6× augmentation operating in the rotated (forward, lateral, log_dt) frame.
|
||||
|
||||
Variants:
|
||||
0 — original
|
||||
1 — lateral flip: lateral → −lateral
|
||||
2 — speed ×0.8: log_Δt[1:] += log(1.25)
|
||||
3 — speed ×1.2: log_Δt[1:] += log(1/1.2)
|
||||
4 — temporal noise: log_Δt[1:] += N(0, 0.05)
|
||||
5 — combined: lateral flip + speed ×0.9
|
||||
|
||||
Args:
|
||||
seq: (N, T, 3) — (forward, lateral, log_dt)
|
||||
cond: (N, 3) — [dist_norm, log_dist, log_dur]
|
||||
|
||||
Returns:
|
||||
seq_aug: (6N, T, 3)
|
||||
cond_aug: (6N, 3)
|
||||
"""
|
||||
log_1_25 = math.log(1.25)
|
||||
log_inv_1_2 = math.log(1.0 / 1.2)
|
||||
log_1_1 = math.log(1.0 / 0.9)
|
||||
|
||||
seqs = [seq]
|
||||
conds = [cond]
|
||||
|
||||
# 1. Lateral flip
|
||||
s1 = seq.copy()
|
||||
s1[:, :, 1] = -s1[:, :, 1]
|
||||
seqs.append(s1)
|
||||
conds.append(cond.copy())
|
||||
|
||||
# 2. Speed ×0.8 (longer duration: log_dt += log(1.25))
|
||||
s2 = seq.copy()
|
||||
s2[:, 1:, 2] += log_1_25
|
||||
c2 = cond.copy()
|
||||
c2[:, 2] += log_1_25 # log_dur updated
|
||||
seqs.append(s2)
|
||||
conds.append(c2)
|
||||
|
||||
# 3. Speed ×1.2 (shorter duration: log_dt += log(1/1.2))
|
||||
s3 = seq.copy()
|
||||
s3[:, 1:, 2] += log_inv_1_2
|
||||
c3 = cond.copy()
|
||||
c3[:, 2] += log_inv_1_2
|
||||
seqs.append(s3)
|
||||
conds.append(c3)
|
||||
|
||||
# 4. Temporal noise
|
||||
s4 = seq.copy()
|
||||
noise = np.random.normal(0.0, 0.05, size=s4[:, 1:, 2].shape).astype(np.float32)
|
||||
s4[:, 1:, 2] += noise
|
||||
seqs.append(s4)
|
||||
conds.append(cond.copy())
|
||||
|
||||
# 5. Lateral flip + speed ×0.9 (i.e. log_dt += log(1/0.9))
|
||||
s5 = seq.copy()
|
||||
s5[:, :, 1] = -s5[:, :, 1]
|
||||
s5[:, 1:, 2] += log_1_1
|
||||
c5 = cond.copy()
|
||||
c5[:, 2] += log_1_1
|
||||
seqs.append(s5)
|
||||
conds.append(c5)
|
||||
|
||||
return np.concatenate(seqs, axis=0), np.concatenate(conds, axis=0)
|
||||
|
||||
|
||||
class TrajectoryDataset(torch.utils.data.Dataset):
|
||||
"""Trajectory dataset with on-the-fly 6× augmentation.
|
||||
|
||||
Replaces the old eager `_augment(seq, cond)` which expanded the dataset
|
||||
6× in memory before training. With this class, the original (N, T, 3)
|
||||
arrays stay as-is and each `__getitem__` call computes one of the 6
|
||||
augmentation variants on demand.
|
||||
|
||||
Augmentation variants (matching legacy `_augment` semantics):
|
||||
0 — original
|
||||
1 — lateral flip (lateral → −lateral)
|
||||
2 — speed ×0.8 (log_dt[1:] += log(1.25), cond[2] += log(1.25))
|
||||
3 — speed ×1.2 (log_dt[1:] += log(1/1.2), cond[2] += log(1/1.2))
|
||||
4 — temporal noise (log_dt[1:] += N(0, 0.05))
|
||||
5 — flip + speed ×0.9 (lateral flip, log_dt[1:] += log(1/0.9), cond[2] += log(1/0.9))
|
||||
"""
|
||||
|
||||
_LOG_1_25 = math.log(1.25)
|
||||
_LOG_INV_1_2 = math.log(1.0 / 1.2)
|
||||
_LOG_INV_0_9 = math.log(1.0 / 0.9)
|
||||
|
||||
def __init__(self, seq: np.ndarray, cond: np.ndarray, augment: bool = True):
|
||||
self.seq = seq
|
||||
self.cond = cond
|
||||
self.augment = augment
|
||||
self._n_aug = 6 if augment else 1
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.seq) * self._n_aug
|
||||
|
||||
def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
base = idx // self._n_aug
|
||||
aug_id = idx % self._n_aug
|
||||
|
||||
s = self.seq[base].copy()
|
||||
c = self.cond[base].copy()
|
||||
|
||||
if aug_id == 1:
|
||||
s[:, 1] = -s[:, 1]
|
||||
elif aug_id == 2:
|
||||
s[1:, 2] += self._LOG_1_25
|
||||
c[2] += self._LOG_1_25
|
||||
elif aug_id == 3:
|
||||
s[1:, 2] += self._LOG_INV_1_2
|
||||
c[2] += self._LOG_INV_1_2
|
||||
elif aug_id == 4:
|
||||
noise = np.random.normal(0.0, 0.05, size=s[1:, 2].shape).astype(np.float32)
|
||||
s[1:, 2] += noise
|
||||
elif aug_id == 5:
|
||||
s[:, 1] = -s[:, 1]
|
||||
s[1:, 2] += self._LOG_INV_0_9
|
||||
c[2] += self._LOG_INV_0_9
|
||||
|
||||
return torch.from_numpy(s), torch.from_numpy(c)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Duration distribution (per distance bin)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _compute_duration_dist(data_path: Path) -> dict:
|
||||
"""Compute per-distance-bin log-normal parameters for trace duration.
|
||||
|
||||
Args:
|
||||
data_path: path to traces.jsonl
|
||||
|
||||
Returns:
|
||||
dict with keys "bins" (list of floats) and "params" (list of dicts
|
||||
with "mu_log" and "sigma_log" per bin).
|
||||
"""
|
||||
bin_durations: list[list[float]] = [[] for _ in range(len(_DIST_BINS) - 1)]
|
||||
|
||||
for raw_line in Path(data_path).read_text(encoding="utf-8").splitlines():
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
trace = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
meta = trace["meta"]
|
||||
events = trace["events"]
|
||||
dist = float(meta.get("dist", 0))
|
||||
|
||||
moves = [e for e in events if e["type"] == "move"]
|
||||
if len(moves) < 2:
|
||||
continue
|
||||
dur = float(moves[-1]["t"] - moves[0]["t"])
|
||||
if dur <= 0:
|
||||
continue
|
||||
|
||||
# Find bin
|
||||
for i in range(len(_DIST_BINS) - 1):
|
||||
if _DIST_BINS[i] <= dist < _DIST_BINS[i + 1]:
|
||||
bin_durations[i].append(dur)
|
||||
break
|
||||
|
||||
params = []
|
||||
for durs in bin_durations:
|
||||
if len(durs) >= 2:
|
||||
log_durs = np.log(np.array(durs, dtype=float))
|
||||
mu_log = float(np.mean(log_durs))
|
||||
sigma_log = float(np.std(log_durs, ddof=1))
|
||||
else:
|
||||
mu_log = float(np.log(500.0))
|
||||
sigma_log = 0.5
|
||||
params.append({"mu_log": mu_log, "sigma_log": max(sigma_log, 0.05)})
|
||||
|
||||
return {"bins": _DIST_BINS, "params": params}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main training function
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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,
|
||||
) -> None:
|
||||
"""Train TrajectoryFlowModel with OT-Conditional Flow Matching.
|
||||
|
||||
Args:
|
||||
data_path: path to traces.jsonl
|
||||
output_dir: directory where artefacts are written
|
||||
epochs: training epochs
|
||||
batch_size: mini-batch size
|
||||
lr: AdamW learning rate
|
||||
seq_len: number of time steps per trajectory
|
||||
progress_callback: optional callable invoked each epoch with
|
||||
{"epoch": n, "total": N, "loss": f}.
|
||||
Called once at the end with {"done": True, "mu", "sigma"}.
|
||||
config: optional TrainConfig for model hyperparameters
|
||||
resume_from: if given, load model weights from this checkpoint
|
||||
directory (must contain flow_model.pt). Used for
|
||||
two-stage training (pretrain → fine-tune).
|
||||
"""
|
||||
data_path = Path(data_path)
|
||||
output_dir = Path(output_dir)
|
||||
|
||||
if config is None:
|
||||
config = TrainConfig(
|
||||
epochs=epochs, batch_size=batch_size, lr=lr, seq_len=seq_len
|
||||
)
|
||||
|
||||
if not data_path.exists():
|
||||
raise FileNotFoundError(f"Data file not found: {data_path}")
|
||||
|
||||
# ---- Load & prepare ----
|
||||
logger.info("Loading data from %s", data_path)
|
||||
seq_np, cond_np, click_durs = load_and_prepare_data(data_path, seq_len=seq_len)
|
||||
logger.info("Loaded %d traces", len(seq_np))
|
||||
|
||||
# ---- Build streaming dataset (on-the-fly 6× augmentation) ----
|
||||
if config.augment:
|
||||
logger.info("Using on-the-fly 6× augmentation, base samples: %d", len(seq_np))
|
||||
ds = TrajectoryDataset(seq_np, cond_np, augment=config.augment)
|
||||
logger.info("Effective dataset size: %d", len(ds))
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# ---- Model & optimiser ----
|
||||
model = TrajectoryFlowModel(
|
||||
seq_len=seq_len,
|
||||
d_model=config.d_model,
|
||||
nhead=config.nhead,
|
||||
num_layers=config.num_layers,
|
||||
dim_feedforward=config.dim_feedforward,
|
||||
dropout=config.dropout,
|
||||
cond_dim=3,
|
||||
)
|
||||
|
||||
# ---- Resume from checkpoint if requested ----
|
||||
if resume_from is not None:
|
||||
resume_path = Path(resume_from) / "flow_model.pt"
|
||||
if not resume_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"resume_from checkpoint not found: {resume_path}"
|
||||
)
|
||||
logger.info("Resuming from checkpoint: %s", resume_path)
|
||||
state_dict = torch.load(resume_path, map_location="cpu", weights_only=True)
|
||||
model.load_state_dict(state_dict)
|
||||
|
||||
optimiser = torch.optim.AdamW(model.parameters(), lr=lr)
|
||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=epochs)
|
||||
|
||||
loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
|
||||
|
||||
# ---- Training loop: OT-Conditional Flow Matching ----
|
||||
model.train()
|
||||
for epoch in range(epochs):
|
||||
epoch_loss = 0.0
|
||||
n_batches = 0
|
||||
|
||||
for x1_batch, cond_batch in loader:
|
||||
B = x1_batch.shape[0]
|
||||
|
||||
# Sample noise x0 ~ N(0, I)
|
||||
x0 = torch.randn_like(x1_batch)
|
||||
|
||||
# Sample random timestep t ~ U[0, 1]
|
||||
t = torch.rand(B)
|
||||
|
||||
# Interpolate: x_t = (1-t)*x0 + t*x1
|
||||
t_expand = t[:, None, None] # (B, 1, 1) for broadcasting
|
||||
x_t = (1.0 - t_expand) * x0 + t_expand * x1_batch
|
||||
|
||||
# Target velocity: v = x1 - x0
|
||||
v_target = x1_batch - x0
|
||||
|
||||
# Predict velocity
|
||||
v_pred = model(x_t, t, cond_batch)
|
||||
|
||||
# MSE loss
|
||||
loss = torch.nn.functional.mse_loss(v_pred, v_target)
|
||||
|
||||
optimiser.zero_grad()
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
||||
optimiser.step()
|
||||
|
||||
epoch_loss += loss.item()
|
||||
n_batches += 1
|
||||
|
||||
scheduler.step()
|
||||
|
||||
epoch_loss /= max(n_batches, 1)
|
||||
logger.debug("Epoch %d/%d loss=%.6f", epoch + 1, epochs, epoch_loss)
|
||||
|
||||
if progress_callback is not None:
|
||||
progress_callback({"epoch": epoch + 1, "total": epochs, "loss": epoch_loss})
|
||||
|
||||
# ---- Save model ----
|
||||
torch.save(model.state_dict(), output_dir / "flow_model.pt")
|
||||
logger.info("Saved flow_model.pt to %s", output_dir)
|
||||
|
||||
# ---- Click duration distribution ----
|
||||
if click_durs:
|
||||
arr = np.array(click_durs)
|
||||
arr = arr[(arr >= 20) & (arr <= 500)]
|
||||
if len(arr) >= 2:
|
||||
mu_c = float(arr.mean())
|
||||
sigma_c = max(float(arr.std()), 1.0)
|
||||
else:
|
||||
mu_c, sigma_c = 80.0, 30.0
|
||||
else:
|
||||
mu_c, sigma_c = 80.0, 30.0
|
||||
|
||||
click_dist = {"mu": mu_c, "sigma": sigma_c, "low": 20.0, "high": 500.0}
|
||||
(output_dir / "click_dist.json").write_text(json.dumps(click_dist, indent=2))
|
||||
|
||||
# ---- Duration distribution (per distance bin) ----
|
||||
dur_dist = _compute_duration_dist(data_path)
|
||||
(output_dir / "duration_dist.json").write_text(json.dumps(dur_dist, indent=2))
|
||||
|
||||
# ---- Train config ----
|
||||
train_cfg = {
|
||||
"seq_len": seq_len,
|
||||
"epochs": epochs,
|
||||
"batch_size": batch_size,
|
||||
"lr": lr,
|
||||
"d_model": config.d_model,
|
||||
"nhead": config.nhead,
|
||||
"num_layers": config.num_layers,
|
||||
"dim_feedforward": config.dim_feedforward,
|
||||
"dropout": config.dropout,
|
||||
"cond_dim": 3,
|
||||
}
|
||||
(output_dir / "train_config.json").write_text(json.dumps(train_cfg, indent=2))
|
||||
|
||||
if progress_callback is not None:
|
||||
progress_callback({"done": True, "mu": mu_c, "sigma": sigma_c})
|
||||
Reference in New Issue
Block a user