"""Train a real DANCE checkpoint on BI2014a for the braindecode tutorial. Reuses the *exact* data pipeline, target builder and model construction from ``examples/applied_examples/plot_dance_event_detection.py`` so the resulting ``model.pt`` (a plain ``state_dict``) loads with ``strict=True`` into the model the tutorial builds. The only differences from the tutorial are training scale: more subjects, more epochs, a OneCycle schedule, minibatches, and keeping the best checkpoint by held-out F1-event. Usage: python train_checkpoint.py --train 1 2 4 5 6 7 8 9 --test 3 --epochs 100 """ import argparse import numpy as np import torch from sklearn.metrics import f1_score from sklearn.preprocessing import robust_scale from torch.utils.data import DataLoader from braindecode.datasets import MOABBDataset from braindecode.models import DANCE from braindecode.preprocessing import ( Preprocessor, create_fixed_length_windows, preprocess, ) from braindecode.training import DanceLoss, f1_event from braindecode.util import set_random_seeds SFREQ = 128.0 WINDOW_S, N_CLASSES, NUM_LATENTS, MAX_EVENTS = 32.0, 3, 256, 150 WINDOW_SAMPLES = int(WINDOW_S * SFREQ) # --- tutorial helpers (verbatim) ------------------------------------------- def robust_scale_clamp(data): return np.clip(robust_scale(data, axis=1), -16, 16) def bi_annotations_to_events(raw): label_to_class = {"NonTarget": 1, "Target": 2} events = [] for ann in raw.annotations: cls = label_to_class.get(str(ann["description"])) if cls is None: continue events.append((float(ann["onset"]), float(ann["onset"] + ann["duration"]), cls)) return events def dance_target_builder(annotations, window_onset, window_duration, max_events, num_latents): start = torch.zeros(max_events) end = torch.zeros(max_events) cls = torch.zeros(max_events, dtype=torch.long) w0, wd = window_onset, window_duration kept = 0 for s, e, c in annotations: s_c, e_c = max(s, w0), min(e, w0 + wd) if e_c <= s_c or int(c) == 0 or kept >= max_events: continue start[kept] = (s_c - w0) / wd end[kept] = (e_c - w0) / wd cls[kept] = int(c) kept += 1 dense = torch.zeros(num_latents, dtype=torch.long) s_tok = (start * num_latents).clamp(0, num_latents).long() e_tok = (end * num_latents).clamp(0, num_latents).long() for i in range(kept): a, b = int(s_tok[i]), int(e_tok[i]) if a < b: dense[a:b] = int(cls[i]) return {"start": start, "end": end, "class": cls, "dense": dense} def dance_collate(batch): eeg = torch.stack([b[0] for b in batch]) out = {"eeg": eeg} for key in ("start", "end", "class", "dense"): out[key] = torch.stack([b[1][key] for b in batch]) return out def detections_to_events(detections, duration): probs = torch.softmax(detections["class"], dim=-1) confidence, label = probs.max(dim=-1) start = detections["start"] * duration end = detections["end"] * duration events = [] for bi in range(label.shape[0]): keep = label[bi] != 0 events.append( list( zip( start[bi, keep].tolist(), end[bi, keep].tolist(), label[bi, keep].tolist(), confidence[bi, keep].tolist(), ) ) ) return events def build_samples(subject_ids): dataset = MOABBDataset(dataset_name="BI2014a", subject_ids=subject_ids) preprocess( dataset, [ Preprocessor("pick_types", eeg=True, stim=False), Preprocessor("filter", l_freq=0.1, h_freq=100.0), Preprocessor("resample", sfreq=SFREQ), Preprocessor(robust_scale_clamp, apply_on_array=True), ], ) windows_ds = create_fixed_length_windows( dataset, window_size_samples=WINDOW_SAMPLES, window_stride_samples=WINDOW_SAMPLES, drop_last_window=True, preload=True, use_mne_epochs=False, ) raw_events = { ds.description["subject"]: bi_annotations_to_events(ds.raw) for ds in windows_ds.datasets } metadata = windows_ds.get_metadata() samples, subjects = [], [] for i in range(len(windows_ds)): x, _, crop_inds = windows_ds[i] eeg = torch.as_tensor(np.asarray(x), dtype=torch.float32) subject = int(metadata.iloc[i]["subject"]) window_onset = float(crop_inds[1]) / SFREQ target = dance_target_builder( raw_events[subject], window_onset, WINDOW_S, MAX_EVENTS, NUM_LATENTS ) samples.append((eeg, target)) subjects.append(subject) chs_info = windows_ds.datasets[0].raw.info["chs"] return samples, np.asarray(subjects), chs_info @torch.no_grad() def evaluate(model, loader, device): model.eval() ev_f1s, dense_preds, dense_targets = [], [], [] for batch in loader: batch = {k: v.to(device) for k, v in batch.items()} out = model.detect(batch["eeg"]) pred_events = detections_to_events(out, duration=WINDOW_S) for bi in range(batch["eeg"].shape[0]): gt = [ (float(s) * WINDOW_S, float(e) * WINDOW_S, int(c)) for s, e, c in zip(batch["start"][bi], batch["end"][bi], batch["class"][bi]) if int(c) != 0 ] preds = [(s, e, c) for (s, e, c, _c) in pred_events[bi]] ev_f1s.append(f1_event(preds, gt, iou_threshold=0.5)) dense_preds.append(out["dense"].argmax(-1).reshape(-1).cpu()) dense_targets.append(batch["dense"].reshape(-1).cpu()) dp = torch.cat(dense_preds).numpy() dt = torch.cat(dense_targets).numpy() sample_f1 = f1_score(dt, dp, labels=list(range(N_CLASSES)), average="macro") return float(np.mean(ev_f1s)), float(sample_f1) def main(): ap = argparse.ArgumentParser() ap.add_argument("--train", type=int, nargs="+", default=[1, 2, 4, 5, 6, 7, 8, 9]) ap.add_argument("--test", type=int, default=3) ap.add_argument("--epochs", type=int, default=100) ap.add_argument("--batch-size", type=int, default=8) ap.add_argument("--max-lr", type=float, default=5e-4) ap.add_argument("--onecycle", action="store_true", help="use OneCycle LR (else constant)") ap.add_argument("--out", type=str, default="/private/home/jarod/dance_ckpt/model.pt") args = ap.parse_args() set_random_seeds(seed=0, cuda=torch.cuda.is_available()) device = "cuda" if torch.cuda.is_available() else "cpu" all_subjects = sorted(set(args.train) | {args.test}) print(f"Loading BI2014a subjects {all_subjects} (test={args.test}) ...", flush=True) samples, subjects, chs_info = build_samples(all_subjects) train_idx = np.flatnonzero(subjects != args.test) test_idx = np.flatnonzero(subjects == args.test) train_samples = [samples[i] for i in train_idx] test_samples = [samples[i] for i in test_idx] print(f"{len(train_samples)} train windows, {len(test_samples)} test windows, " f"n_chans={len(chs_info)}", flush=True) train_loader = DataLoader(train_samples, batch_size=args.batch_size, shuffle=True, collate_fn=dance_collate, drop_last=True) test_loader = DataLoader(test_samples, batch_size=len(test_samples), collate_fn=dance_collate) model = DANCE( n_outputs=N_CLASSES, n_chans=len(chs_info), chs_info=chs_info, n_times=WINDOW_SAMPLES, sfreq=SFREQ, input_window_seconds=WINDOW_S, ).to(device) criterion = DanceLoss(num_latents=NUM_LATENTS) optimizer = torch.optim.Adam(model.parameters(), lr=args.max_lr) steps = max(1, len(train_loader)) if args.onecycle: sched = torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr=args.max_lr, total_steps=args.epochs * steps, pct_start=0.1 ) else: sched = None best_f1, best_state = -1.0, None for epoch in range(args.epochs): model.train() ep_loss = 0.0 for batch in train_loader: batch = {k: v.to(device) for k, v in batch.items()} optimizer.zero_grad() out = model.detect(batch["eeg"]) loss, _ = criterion(out, batch, duration=WINDOW_S) loss.backward() optimizer.step() if sched is not None: sched.step() ep_loss += float(loss) ev_f1, samp_f1 = evaluate(model, test_loader, device) if ev_f1 > best_f1: best_f1 = ev_f1 best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} torch.save(best_state, args.out) if epoch % 5 == 0 or epoch == args.epochs - 1: print(f"ep{epoch:03d} loss={ep_loss/steps:.3f} " f"F1-event={ev_f1:.3f} F1-sample={samp_f1:.3f} (best={best_f1:.3f})", flush=True) print(f"\nBEST held-out F1-event={best_f1:.3f}; checkpoint saved to {args.out}", flush=True) if __name__ == "__main__": main()