train_checkpoint.py

"""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()