From c3ac48278ed1d3669da5ab8127deeb2ebf8232a1 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 18 May 2026 21:27:02 +0000 Subject: [PATCH] =?UTF-8?q?[fix]=20train=20on=20trial-aligned=20windows=20?= =?UTF-8?q?for=20correct=20R=C2=B2=20evaluation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch train_benchmark.py from NLBDataset (sliding windows over the full recording) to TrialAlignedDataset (one window per trial, -100ms to +500ms around move_onset_time). The sliding-window approach samples ~85% rest periods, collapsing R² to ~0; trial-aligned windows match the published NLB evaluation protocol and yield meaningful R² (Wiener baseline ~0.49). https://claude.ai/code/session_01AcwBjnbswWKkeyujGsuFzy --- scripts/train_benchmark.py | 43 +++++++++++++++----------------------- 1 file changed, 17 insertions(+), 26 deletions(-) diff --git a/scripts/train_benchmark.py b/scripts/train_benchmark.py index ed7f0be..6d680c4 100644 --- a/scripts/train_benchmark.py +++ b/scripts/train_benchmark.py @@ -19,7 +19,7 @@ import torch from torch.utils.data import DataLoader -from cortex.data.nlb import NLBDataset, collate_events +from cortex.data.nlb import build_trial_aligned_datasets, collate_events from cortex.models import CortexModel from cortex.models.config import CORTEX_S from cortex.training.eval import evaluate, r2_score @@ -36,29 +36,16 @@ def build_loaders( - data_root: str, batch_size: int, val_max_samples: int = 2048 + nwb_path: str, batch_size: int ) -> tuple[DataLoader, DataLoader]: - kwargs = dict( - data_root=data_root, - dandiset_id="000128", + train_ds, val_ds = build_trial_aligned_datasets( + nwb_path=Path(nwb_path), bin_size_ms=5, - window_ms=600, - stride_ms=50, + pre_onset_ms=100, + post_onset_ms=500, max_neurons=CORTEX_S.max_neurons, spike_value_buckets=CORTEX_S.spike_value_buckets, - download=False, ) - train_ds = NLBDataset(split="train", **kwargs) # type: ignore[arg-type] - val_ds = NLBDataset(split="val", **kwargs) # type: ignore[arg-type] - - # Cap val at val_max_samples to keep eval fast during training - if len(val_ds) > val_max_samples: - from torch.utils.data import Subset - import random - rng = random.Random(42) - val_idx = rng.sample(range(len(val_ds)), val_max_samples) - val_ds = Subset(val_ds, val_idx) # type: ignore[assignment] - loader_kw = dict(batch_size=batch_size, num_workers=0, collate_fn=collate_events) return ( DataLoader(train_ds, shuffle=True, **loader_kw), @@ -73,18 +60,22 @@ def main() -> None: parser.add_argument("--batch-size", type=int, default=32) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--warmup-steps", type=int, default=200) - parser.add_argument("--data-root", default="./data") + parser.add_argument( + "--nwb-path", + default="./data/000128/sub-Jenkins/sub-Jenkins_ses-full_desc-train_behavior+ecephys.nwb", + help="Path to the MC_Maze NWB file (train+behavior variant).", + ) parser.add_argument("--out", default="benchmarks/training/results_raw.json") args = parser.parse_args() device = select_device(preference=args.device) log.info("benchmark_start", device=str(device), max_steps=args.max_steps, model_params=f"{sum(p.numel() for p in CortexModel(CORTEX_S).parameters())/1e6:.1f}M") - log.info("loading_data") + log.info("loading_data", nwb_path=args.nwb_path) t0 = time.time() - train_loader, val_loader = build_loaders(args.data_root, args.batch_size) - log.info("data_loaded", train_windows=len(train_loader.dataset), # type: ignore[arg-type] - val_windows=len(val_loader.dataset), elapsed_s=f"{time.time()-t0:.1f}") # type: ignore[arg-type] + train_loader, val_loader = build_loaders(args.nwb_path, args.batch_size) + log.info("data_loaded", train_trials=len(train_loader.dataset), # type: ignore[arg-type] + val_trials=len(val_loader.dataset), elapsed_s=f"{time.time()-t0:.1f}") # type: ignore[arg-type] model = CortexModel(CORTEX_S).to(device) optimizer = build_optimizer(model, _cfg(args.lr)) @@ -159,8 +150,8 @@ def main() -> None: "batch_size": args.batch_size, "final_r2_velocity": eval_log[-1]["r2_velocity"] if eval_log else None, "final_mse_velocity": eval_log[-1]["mse_velocity"] if eval_log else None, - "train_windows": len(train_loader.dataset), # type: ignore[arg-type] - "val_windows": len(val_loader.dataset), # type: ignore[arg-type] + "train_trials": len(train_loader.dataset), # type: ignore[arg-type] + "val_trials": len(val_loader.dataset), # type: ignore[arg-type] "total_elapsed_s": round(total_elapsed, 1), "avg_step_ms": round(avg_step_ms, 1), "p99_step_ms": round(p99_step_ms, 1),