From bba8fc567b98fc2773e9f3adbb122abfe7d05342 Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Sun, 10 May 2026 12:51:52 +0800 Subject: [PATCH] feat(trainer): replace eager _augment with streaming TrajectoryDataset Co-Authored-By: Claude Sonnet 4.6 --- ai_mouse/trainer.py | 70 ++++++++++++++++++++++++++++++++++++++----- tests/test_trainer.py | 51 +++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 8 deletions(-) diff --git a/ai_mouse/trainer.py b/ai_mouse/trainer.py index 3360fe5..9e5864c 100644 --- a/ai_mouse/trainer.py +++ b/ai_mouse/trainer.py @@ -21,7 +21,7 @@ from pathlib import Path import numpy as np import torch -from torch.utils.data import DataLoader, TensorDataset +from torch.utils.data import DataLoader from ai_mouse.config import TrainConfig from ai_mouse.coord import encode_trajectory @@ -226,6 +226,62 @@ def _augment( 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_1_1 = 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_1_1 + c[2] += self._LOG_1_1 + + return torch.from_numpy(s), torch.from_numpy(c) + + # --------------------------------------------------------------------------- # Duration distribution (per distance bin) # --------------------------------------------------------------------------- @@ -328,12 +384,11 @@ def train( seq_np, cond_np, click_durs = load_and_prepare_data(data_path, seq_len=seq_len) logger.info("Loaded %d traces", len(seq_np)) - # ---- Augment ---- - seq_np, cond_np = _augment(seq_np, cond_np) - logger.info("After augmentation: %d samples", len(seq_np)) - - seq_t = torch.from_numpy(seq_np) # (N, seq_len, 3) - cond_t = torch.from_numpy(cond_np) # (N, 3) + # ---- 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) @@ -350,7 +405,6 @@ def train( optimiser = torch.optim.AdamW(model.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=epochs) - ds = TensorDataset(seq_t, cond_t) loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False) # ---- Training loop: OT-Conditional Flow Matching ---- diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 5af97c6..3e61c21 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -119,3 +119,54 @@ class TestTrain: first_half = np.mean(losses[:10]) second_half = np.mean(losses[10:]) assert second_half < first_half + + +class TestTrajectoryDataset: + def test_dataset_length_with_augmentation(self): + """Dataset length = N * 6 when augment=True.""" + from ai_mouse.trainer import TrajectoryDataset + seq = np.zeros((10, 64, 3), dtype=np.float32) + cond = np.zeros((10, 3), dtype=np.float32) + ds = TrajectoryDataset(seq, cond, augment=True) + assert len(ds) == 60 + + def test_dataset_length_without_augmentation(self): + from ai_mouse.trainer import TrajectoryDataset + seq = np.zeros((10, 64, 3), dtype=np.float32) + cond = np.zeros((10, 3), dtype=np.float32) + ds = TrajectoryDataset(seq, cond, augment=False) + assert len(ds) == 10 + + def test_getitem_returns_tensors(self): + from ai_mouse.trainer import TrajectoryDataset + import torch + seq = np.random.randn(5, 64, 3).astype(np.float32) + cond = np.random.randn(5, 3).astype(np.float32) + ds = TrajectoryDataset(seq, cond, augment=True) + s, c = ds[0] + assert isinstance(s, torch.Tensor) + assert isinstance(c, torch.Tensor) + assert s.shape == (64, 3) + assert c.shape == (3,) + + def test_aug_id_zero_returns_original(self): + """Aug id 0 (idx=0 % 6 == 0) should return the original sample unchanged.""" + from ai_mouse.trainer import TrajectoryDataset + import torch + seq = np.array([[[0.5, 0.7, 0.3]] * 64] * 3, dtype=np.float32) + cond = np.array([[1.0, 2.0, 3.0]] * 3, dtype=np.float32) + ds = TrajectoryDataset(seq, cond, augment=True) + s0, c0 = ds[0] + np.testing.assert_allclose(s0.numpy(), seq[0], rtol=1e-5) + np.testing.assert_allclose(c0.numpy(), cond[0], rtol=1e-5) + + def test_aug_id_one_flips_lateral(self): + """Aug id 1 should flip the sign of the lateral channel (index 1).""" + from ai_mouse.trainer import TrajectoryDataset + seq = np.zeros((1, 64, 3), dtype=np.float32) + seq[0, :, 1] = 0.5 # lateral all positive + cond = np.zeros((1, 3), dtype=np.float32) + ds = TrajectoryDataset(seq, cond, augment=True) + # idx=1 → base_idx=0, aug_id=1 → flip + s1, _ = ds[1] + assert (s1[:, 1] < 0).all()