Upload train_checkpoint.py with huggingface_hub

train_checkpoint.py

@@ -0,0 +1,240 @@
+"""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()