chore: import upstream snapshot with attribution

2026-07-13 11:59:54 +08:00
commit 9740bc64c9
2636 changed files with 1330642 additions and 0 deletions
@@ -0,0 +1,761 @@
+#!/usr/bin/env python3
+"""Train the person-count head — ADR-103 v0.0.1.
+
+Mirrors the Conv1d encoder architecture from cog-person-count's
+`src/inference.rs::CountNet` exactly, so the learned weights load
+into the Rust cog without translation. Trains on
+data/paired/wiflow-p7-1779210883.paired.jsonl (1,077 samples with
+n_persons_mode labels in {0, 1}).
+
+Output: count_v1.safetensors + count_v1.onnx + train_results.json.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import struct
+import time
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Architecture constants — MUST match cog-person-count's src/inference.rs.
+N_SUB = 56
+N_FRAMES = 20
+COUNT_CLASSES = 8
+
+
+class CountNet(nn.Module):
+    """Mirrors cog_person_count::inference::CountNet bit-for-bit."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        # Encoder — identical to the pose cog's encoder so future joint
+        # training can share weights.
+        self.enc_c1 = nn.Conv1d(N_SUB, 64, kernel_size=3, padding=1, dilation=1)
+        self.enc_c2 = nn.Conv1d(64, 128, kernel_size=3, padding=2, dilation=2)
+        self.enc_c3 = nn.Conv1d(128, 128, kernel_size=3, padding=4, dilation=4)
+        # Count head
+        self.count_head_fc1 = nn.Linear(128, 64)
+        self.count_head_fc2 = nn.Linear(64, COUNT_CLASSES)
+        # Confidence head
+        self.conf_head_fc1 = nn.Linear(128, 32)
+        self.conf_head_fc2 = nn.Linear(32, 1)
+
+    def forward(self, x: torch.Tensor):
+        # x: [B, 56, 20]
+        h = F.relu(self.enc_c1(x))
+        h = F.relu(self.enc_c2(h))
+        h = F.relu(self.enc_c3(h))
+        h = h.mean(dim=2)  # [B, 128]
+
+        # Logits (un-normalised); softmax at inference + cross-entropy training.
+        c = F.relu(self.count_head_fc1(h))
+        count_logits = self.count_head_fc2(c)
+
+        # Confidence head — sigmoid at inference; BCE-with-logits at training.
+        cf = F.relu(self.conf_head_fc1(h))
+        conf_logits = self.conf_head_fc2(cf)
+
+        return count_logits, conf_logits
+
+
+def load_paired(path: Path) -> tuple[np.ndarray, np.ndarray]:
+    """Return (X, y) where X is [N, 56, 20] CSI and y is [N] integer counts."""
+    csis, ys = [], []
+    with path.open(encoding="utf-8") as f:
+        for line in f:
+            if not line.strip():
+                continue
+            d = json.loads(line)
+            shape = d.get("csi_shape", [N_SUB, N_FRAMES])
+            if shape != [N_SUB, N_FRAMES]:
+                continue
+            csi = np.asarray(d["csi"], dtype=np.float32).reshape(N_SUB, N_FRAMES)
+            csis.append(csi)
+            ys.append(int(d.get("n_persons_mode", 0)))
+    X = np.stack(csis, axis=0)
+    y = np.asarray(ys, dtype=np.int64)
+    return X, y
+
+
+def temporal_split(X: np.ndarray, y: np.ndarray, eval_frac: float = 0.2):
+    """Held-out time-window eval (last `eval_frac` of samples, by index)."""
+    n = X.shape[0]
+    n_eval = int(round(n * eval_frac))
+    n_train = n - n_eval
+    return (
+        X[:n_train], y[:n_train],
+        X[n_train:], y[n_train:],
+    )
+
+
+def stratified_k_fold(X: np.ndarray, y: np.ndarray, k: int = 5):
+    """Stratified k-fold cross-validation splits — hand-rolled, no sklearn.
+
+    Per class: shuffle the indices (deterministic seed 42), split into k
+    near-equal chunks, then assemble fold i by taking chunk i from every
+    class. Yields (X_train, y_train, X_val, y_val) per fold, with class
+    distribution preserved within ±1.
+    """
+    rng = np.random.default_rng(seed=42)
+    classes = np.unique(y)
+    per_class_folds = {}
+    for c in classes:
+        idx = np.where(y == c)[0]
+        rng.shuffle(idx)
+        per_class_folds[c] = np.array_split(idx, k)
+    for fold in range(k):
+        val_idx = np.concatenate([per_class_folds[c][fold] for c in classes])
+        train_idx = np.concatenate(
+            [per_class_folds[c][f] for c in classes for f in range(k) if f != fold]
+        )
+        yield X[train_idx], y[train_idx], X[val_idx], y[val_idx]
+
+
+def standardise(X_train: np.ndarray, X_eval: np.ndarray):
+    """Z-score by subcarrier across the time axis. Eval uses train stats."""
+    mu = X_train.mean(axis=(0, 2), keepdims=True)
+    sd = X_train.std(axis=(0, 2), keepdims=True) + 1e-6
+    return (X_train - mu) / sd, (X_eval - mu) / sd
+
+
+def write_safetensors(model: CountNet, path: Path):
+    """Write the model's state in the same on-disk layout the Rust cog expects."""
+    state = model.state_dict()
+    # Map PyTorch param names → cog-person-count's VarBuilder paths.
+    rename = {
+        "enc_c1.weight": "enc.c1.weight",
+        "enc_c1.bias":   "enc.c1.bias",
+        "enc_c2.weight": "enc.c2.weight",
+        "enc_c2.bias":   "enc.c2.bias",
+        "enc_c3.weight": "enc.c3.weight",
+        "enc_c3.bias":   "enc.c3.bias",
+        "count_head_fc1.weight": "count_head.fc1.weight",
+        "count_head_fc1.bias":   "count_head.fc1.bias",
+        "count_head_fc2.weight": "count_head.fc2.weight",
+        "count_head_fc2.bias":   "count_head.fc2.bias",
+        "conf_head_fc1.weight":  "conf_head.fc1.weight",
+        "conf_head_fc1.bias":    "conf_head.fc1.bias",
+        "conf_head_fc2.weight":  "conf_head.fc2.weight",
+        "conf_head_fc2.bias":    "conf_head.fc2.bias",
+    }
+
+    header = {}
+    payload = bytearray()
+    offset = 0
+    for torch_name, cog_name in rename.items():
+        t = state[torch_name].detach().cpu().numpy().astype(np.float32)
+        n_bytes = t.nbytes
+        header[cog_name] = {
+            "dtype": "F32",
+            "shape": list(t.shape),
+            "data_offsets": [offset, offset + n_bytes],
+        }
+        payload.extend(t.tobytes())
+        offset += n_bytes
+
+    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    with path.open("wb") as f:
+        f.write(struct.pack("<Q", len(header_bytes)))
+        f.write(header_bytes)
+        f.write(payload)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--paired", required=True)
+    parser.add_argument("--out-safetensors", default="count_v1.safetensors")
+    parser.add_argument("--out-onnx", default="count_v1.onnx")
+    parser.add_argument("--out-results", default="count_train_results.json")
+    parser.add_argument("--epochs", type=int, default=400)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--weight-decay", type=float, default=0.01)
+    parser.add_argument("--k-fold", type=int, default=None, help="If set, run k-fold CV; else use temporal split")
+    parser.add_argument("--v2", action="store_true",
+                        help="v0.0.2 training: random 80/20 split + label smoothing + early stopping "
+                             "+ balanced sampling + temperature-scaled confidence head.")
+    parser.add_argument("--label-smoothing", type=float, default=0.1)
+    parser.add_argument("--patience", type=int, default=20)
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"device: {device}")
+
+    X, y = load_paired(Path(args.paired))
+    print(f"loaded {X.shape[0]} samples, X shape {X.shape}, "
+          f"label distribution: {dict(Counter(y.tolist()).most_common())}")
+
+    # K-fold cross-validation mode
+    if args.k_fold is not None:
+        print(f"\n=== {args.k_fold}-fold cross-validation ===")
+        fold_results = []
+        overall_t0 = time.perf_counter()
+
+        for fold_idx, (X_train, y_train, X_val, y_val) in enumerate(stratified_k_fold(X, y, k=args.k_fold)):
+            print(f"\nFold {fold_idx + 1}/{args.k_fold}")
+            X_train, X_val = standardise(X_train, X_val)
+
+            cls_counts = np.bincount(y_train, minlength=COUNT_CLASSES).astype(np.float32)
+            cls_counts = np.where(cls_counts > 0, cls_counts, 1.0)
+            cls_weight = (1.0 / cls_counts) / (1.0 / cls_counts).sum() * COUNT_CLASSES
+            cls_weight_t = torch.from_numpy(cls_weight).to(device)
+
+            Xt = torch.from_numpy(X_train).to(device)
+            yt = torch.from_numpy(y_train).to(device)
+            Xv = torch.from_numpy(X_val).to(device)
+            yv = torch.from_numpy(y_val).to(device)
+
+            model = CountNet().to(device)
+            opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+            sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=50, T_mult=1)
+
+            n_train = X_train.shape[0]
+            best_eval_acc = 0.0
+            best_state = None
+
+            for epoch in range(args.epochs):
+                model.train()
+                perm = torch.randperm(n_train, device=device)
+                train_loss = 0.0
+                train_correct = 0
+                n_batches = 0
+                for i in range(0, n_train, args.batch_size):
+                    idx = perm[i : i + args.batch_size]
+                    xb = Xt[idx]
+                    yb = yt[idx]
+                    opt.zero_grad()
+                    count_logits, conf_logits = model(xb)
+                    ce = F.cross_entropy(count_logits, yb, weight=cls_weight_t)
+                    with torch.no_grad():
+                        pred = count_logits.argmax(dim=1)
+                        correct_indicator = (pred == yb).float().unsqueeze(1)
+                    bce = F.binary_cross_entropy_with_logits(conf_logits, correct_indicator)
+                    with torch.no_grad():
+                        conf_sigm = torch.sigmoid(conf_logits)
+                    brier = ((conf_sigm - correct_indicator) ** 2).mean()
+                    loss = ce + 0.3 * bce + 0.1 * brier
+                    loss.backward()
+                    opt.step()
+                    train_loss += loss.item()
+                    train_correct += (pred == yb).sum().item()
+                    n_batches += 1
+
+                sched.step()
+
+                model.eval()
+                with torch.no_grad():
+                    cl_v, _ = model(Xv)
+                    eval_pred = cl_v.argmax(dim=1)
+                    eval_acc = (eval_pred == yv).float().mean().item()
+
+                if eval_acc > best_eval_acc:
+                    best_eval_acc = eval_acc
+                    best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
+
+            # Restore best checkpoint and final eval
+            if best_state is not None:
+                model.load_state_dict(best_state)
+
+            model.eval()
+            with torch.no_grad():
+                cl_v, conf_v = model(Xv)
+                pred_v = cl_v.argmax(dim=1)
+                acc = (pred_v == yv).float().mean().item()
+                within1 = ((pred_v - yv).abs() <= 1).float().mean().item()
+                mae = (pred_v - yv).abs().float().mean().item()
+
+                # Per-class accuracy
+                per_class = {}
+                for k in range(COUNT_CLASSES):
+                    mask = yv == k
+                    n = mask.sum().item()
+                    if n > 0:
+                        per_class[k] = {
+                            "support": int(n),
+                            "accuracy": ((pred_v == yv) & mask).sum().item() / n,
+                        }
+
+                # Spearman
+                conf_sigm = torch.sigmoid(conf_v).squeeze(-1)
+                correct = (pred_v == yv).float()
+                c_rank = conf_sigm.argsort().argsort().float()
+                r_rank = correct.argsort().argsort().float()
+                c_centered = c_rank - c_rank.mean()
+                r_centered = r_rank - r_rank.mean()
+                denom = (c_centered.norm() * r_centered.norm()).item()
+                spearman = (c_centered * r_centered).sum().item() / denom if denom > 0 else 0.0
+
+            fold_results.append({
+                "fold": fold_idx + 1,
+                "accuracy": acc,
+                "within_pm1": within1,
+                "mae": mae,
+                "spearman": spearman,
+                "per_class_accuracy": per_class,
+            })
+            print(f"  accuracy={acc:.3f}  within±1={within1:.3f}  mae={mae:.3f}  spearman={spearman:.3f}")
+
+        # K-fold summary
+        total_time = time.perf_counter() - overall_t0
+        accs = [r["accuracy"] for r in fold_results]
+        within1s = [r["within_pm1"] for r in fold_results]
+        maes = [r["mae"] for r in fold_results]
+        spears = [r["spearman"] for r in fold_results]
+
+        print(f"\n=== {args.k_fold}-fold summary ({total_time:.1f} s) ===")
+        print(f"  accuracy:       {np.mean(accs):.3f} ± {np.std(accs):.3f}")
+        print(f"  within ±1:      {np.mean(within1s):.3f} ± {np.std(within1s):.3f}")
+        print(f"  MAE:            {np.mean(maes):.3f} ± {np.std(maes):.3f}")
+        print(f"  conf↔correct Spearman: {np.mean(spears):.3f} ± {np.std(spears):.3f}")
+
+        # Per-class summary across folds
+        for k in range(COUNT_CLASSES):
+            accs_k = [r["per_class_accuracy"].get(k, {}).get("accuracy", 0.0) for r in fold_results]
+            n_k = [r["per_class_accuracy"].get(k, {}).get("support", 0) for r in fold_results]
+            if any(n > 0 for n in n_k):
+                print(f"  class {k}:  {np.mean(accs_k):.3f} mean accuracy (support: {n_k})")
+
+        # Write k-fold results to JSON
+        results = {
+            "mode": "k_fold_cv",
+            "k": args.k_fold,
+            "backend": "pytorch-cuda" if device.type == "cuda" else "pytorch-cpu",
+            "total_time_s": total_time,
+            "fold_results": fold_results,
+            "summary": {
+                "mean_accuracy": float(np.mean(accs)),
+                "std_accuracy": float(np.std(accs)),
+                "mean_within_pm1": float(np.mean(within1s)),
+                "std_within_pm1": float(np.std(within1s)),
+                "mean_mae": float(np.mean(maes)),
+                "std_mae": float(np.std(maes)),
+                "mean_spearman": float(np.mean(spears)),
+                "std_spearman": float(np.std(spears)),
+            },
+            "hyperparameters": {
+                "optimizer": "AdamW",
+                "lr": args.lr,
+                "weight_decay": args.weight_decay,
+                "batch_size": args.batch_size,
+                "schedule": "cosine_warm_restarts",
+                "epochs": args.epochs,
+            },
+        }
+        Path(args.out_results).write_text(json.dumps(results, indent=2))
+        print(f"\nwrote {args.out_results}")
+        return
+
+    # ---------------------------------------------------------------
+    # v0.0.2 training path: random 80/20 + label smoothing + early
+    # stopping + class-balanced batch sampling + temperature scaling.
+    # ---------------------------------------------------------------
+    if args.v2:
+        rng = np.random.default_rng(seed=42)
+        idx = np.arange(X.shape[0])
+        rng.shuffle(idx)
+        n_eval = int(round(0.2 * X.shape[0]))
+        eval_idx, train_idx = idx[:n_eval], idx[n_eval:]
+        X_train, X_eval = X[train_idx], X[eval_idx]
+        y_train, y_eval = y[train_idx], y[eval_idx]
+        X_train, X_eval = standardise(X_train, X_eval)
+        print(f"v0.0.2 mode — random 80/20 split: train={len(y_train)} eval={len(y_eval)}")
+        print(f"  train class dist: {dict(Counter(y_train.tolist()).most_common())}")
+        print(f"  eval  class dist: {dict(Counter(y_eval.tolist()).most_common())}")
+
+        Xt = torch.from_numpy(X_train).to(device)
+        yt = torch.from_numpy(y_train).to(device)
+        Xe = torch.from_numpy(X_eval).to(device)
+        ye = torch.from_numpy(y_eval).to(device)
+
+        # Class-balanced sampler: for each batch, sample with replacement
+        # so each class has equal expected count regardless of dataset
+        # distribution. With our ~533/544 split this is nearly a no-op
+        # but it generalises to imbalanced multi-room data later.
+        cls_counts = np.bincount(y_train, minlength=COUNT_CLASSES).astype(np.float32)
+        cls_counts = np.where(cls_counts > 0, cls_counts, 1.0)
+        per_sample_weight = (1.0 / cls_counts[y_train])
+        per_sample_weight_t = torch.from_numpy(per_sample_weight.astype(np.float32)).to(device)
+
+        model = CountNet().to(device)
+        opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+        sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=50, T_mult=1)
+
+        n_train = X_train.shape[0]
+        batches_per_epoch = max(1, n_train // args.batch_size)
+        epoch_losses = []
+        t0 = time.perf_counter()
+        best_eval_acc = 0.0
+        best_state = None
+        epochs_without_improvement = 0
+
+        for epoch in range(args.epochs):
+            model.train()
+            train_loss = 0.0; train_correct = 0; n_batches = 0
+            for _ in range(batches_per_epoch):
+                # Balanced sample with replacement
+                idx_t = torch.multinomial(per_sample_weight_t, args.batch_size, replacement=True)
+                xb = Xt[idx_t]; yb = yt[idx_t]
+                opt.zero_grad()
+                count_logits, conf_logits = model(xb)
+                ce = F.cross_entropy(count_logits, yb, label_smoothing=args.label_smoothing)
+                with torch.no_grad():
+                    pred = count_logits.argmax(dim=1)
+                    correct_indicator = (pred == yb).float().unsqueeze(1)
+                bce = F.binary_cross_entropy_with_logits(conf_logits, correct_indicator)
+                with torch.no_grad():
+                    conf_sigm = torch.sigmoid(conf_logits)
+                brier = ((conf_sigm - correct_indicator) ** 2).mean()
+                loss = ce + 0.3 * bce + 0.1 * brier
+                loss.backward()
+                opt.step()
+                train_loss += loss.item()
+                train_correct += (pred == yb).sum().item()
+                n_batches += 1
+            sched.step()
+
+            model.eval()
+            with torch.no_grad():
+                cl_e, _ = model(Xe)
+                eval_loss = F.cross_entropy(cl_e, ye).item()
+                eval_pred = cl_e.argmax(dim=1)
+                eval_acc = (eval_pred == ye).float().mean().item()
+            epoch_losses.append({
+                "epoch": epoch,
+                "train_loss": train_loss / max(1, n_batches),
+                "train_acc": train_correct / max(1, n_batches * args.batch_size),
+                "eval_loss": eval_loss,
+                "eval_acc": eval_acc,
+            })
+            if eval_acc > best_eval_acc:
+                best_eval_acc = eval_acc
+                best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
+                epochs_without_improvement = 0
+            else:
+                epochs_without_improvement += 1
+
+            if epoch < 5 or epoch % 25 == 0:
+                print(f"epoch {epoch:3d}  train_loss={train_loss/n_batches:.4f}  "
+                      f"train_acc={train_correct/(n_batches*args.batch_size):.3f}  "
+                      f"eval_loss={eval_loss:.4f}  eval_acc={eval_acc:.3f}  "
+                      f"epochs_no_improve={epochs_without_improvement}")
+            if epochs_without_improvement >= args.patience:
+                print(f"early stopping at epoch {epoch} (no improvement for {args.patience} epochs)")
+                break
+
+        train_time = time.perf_counter() - t0
+        print(f"\ntrained {epoch + 1} epochs in {train_time:.1f} s  (best eval_acc {best_eval_acc:.3f})")
+        if best_state is not None:
+            model.load_state_dict(best_state)
+
+        # Temperature scaling on the confidence head — fit a scalar T s.t.
+        # sigmoid(conf_logits / T) is best-calibrated on the eval set.
+        model.eval()
+        with torch.no_grad():
+            cl_e, conf_e = model(Xe)
+            pred_e = cl_e.argmax(dim=1)
+            correct_indicator = (pred_e == ye).float()
+        # 1D optimisation over T via LBFGS.
+        T = torch.nn.Parameter(torch.ones(1, device=device))
+        opt_t = torch.optim.LBFGS([T], lr=0.1, max_iter=50)
+        def eval_t():
+            opt_t.zero_grad()
+            scaled = conf_e.squeeze(-1) / T
+            loss_t = F.binary_cross_entropy_with_logits(scaled, correct_indicator)
+            loss_t.backward()
+            return loss_t
+        opt_t.step(eval_t)
+        T_val = float(T.detach().cpu().item())
+        print(f"  temperature scale T = {T_val:.4f}")
+
+        # Final eval with temperature applied.
+        with torch.no_grad():
+            cl_e, conf_e = model(Xe)
+            probs_e = F.softmax(cl_e, dim=1)
+            pred_e = cl_e.argmax(dim=1)
+            acc = (pred_e == ye).float().mean().item()
+            within1 = ((pred_e - ye).abs() <= 1).float().mean().item()
+            mae = (pred_e - ye).abs().float().mean().item()
+            per_class = {}
+            for k in range(COUNT_CLASSES):
+                mask = ye == k
+                n = mask.sum().item()
+                if n > 0:
+                    per_class[k] = {
+                        "support": int(n),
+                        "accuracy": ((pred_e == ye) & mask).sum().item() / n,
+                    }
+            conf_sigm = torch.sigmoid(conf_e.squeeze(-1) / T_val)
+            correct = (pred_e == ye).float()
+            c_rank = conf_sigm.argsort().argsort().float()
+            r_rank = correct.argsort().argsort().float()
+            c_centered = c_rank - c_rank.mean()
+            r_centered = r_rank - r_rank.mean()
+            denom = (c_centered.norm() * r_centered.norm()).item()
+            spearman = (c_centered * r_centered).sum().item() / denom if denom > 0 else 0.0
+
+        print(f"\n=== v0.0.2 final eval ===")
+        print(f"  accuracy:       {acc:.3f}")
+        print(f"  within ±1:      {within1:.3f}")
+        print(f"  MAE:            {mae:.3f}")
+        print(f"  conf↔correct Spearman (post-temp): {spearman:.3f}")
+        for k, v in per_class.items():
+            print(f"  class {k}:  {v['accuracy']:.3f} accuracy on {v['support']} samples")
+
+        write_safetensors(model, Path(args.out_safetensors))
+        # Also append the temperature scalar so the cog can apply it.
+        # We add it by appending to the safetensors file using the
+        # write_safetensors helper but with the temperature recorded
+        # as a separate file alongside (count_v1.temperature.txt) for
+        # consumption by the Rust cog inference path.
+        Path(args.out_safetensors + ".temperature").write_text(f"{T_val}\n")
+        print(f"wrote {args.out_safetensors} ({Path(args.out_safetensors).stat().st_size} bytes)")
+        print(f"wrote {args.out_safetensors}.temperature ({T_val})")
+
+        # ONNX
+        dummy = torch.zeros(1, N_SUB, N_FRAMES, device=device)
+        try:
+            torch.onnx.export(model, dummy, args.out_onnx, opset_version=18,
+                              input_names=["csi_window"],
+                              output_names=["count_logits", "conf_logits"],
+                              dynamic_axes={"csi_window": {0: "batch"},
+                                            "count_logits": {0: "batch"},
+                                            "conf_logits": {0: "batch"}},
+                              export_params=True, do_constant_folding=True)
+            print(f"wrote {args.out_onnx} ({Path(args.out_onnx).stat().st_size} bytes)")
+        except Exception as e:
+            print(f"WARN: ONNX export failed: {e}")
+
+        results = {
+            "mode": "v0.0.2",
+            "backend": "pytorch-cuda" if device.type == "cuda" else "pytorch-cpu",
+            "epochs_trained": epoch + 1,
+            "train_time_s": train_time,
+            "best_eval_acc": best_eval_acc,
+            "final_eval_acc": acc,
+            "final_eval_within_pm1": within1,
+            "final_eval_mae": mae,
+            "temperature_scale": T_val,
+            "conf_correctness_spearman_post_temp": spearman,
+            "per_class_accuracy": per_class,
+            "hyperparameters": {
+                "optimizer": "AdamW",
+                "lr": args.lr,
+                "weight_decay": args.weight_decay,
+                "batch_size": args.batch_size,
+                "schedule": "cosine_warm_restarts",
+                "epochs_max": args.epochs,
+                "label_smoothing": args.label_smoothing,
+                "patience": args.patience,
+                "split": "random_80_20_seed_42",
+                "balanced_sampler": True,
+                "temperature_scaling": True,
+            },
+            "epoch_losses": epoch_losses,
+        }
+        Path(args.out_results).write_text(json.dumps(results, indent=2))
+        print(f"wrote {args.out_results}")
+        return
+
+    # Original temporal-split mode (kept for v0.0.1 reproducibility).
+    X_train, y_train, X_eval, y_eval = temporal_split(X, y, eval_frac=0.2)
+    X_train, X_eval = standardise(X_train, X_eval)
+
+    # Re-balance via class weights — handles the 50/50 split fine
+    # but also makes the loss correct under future imbalanced data.
+    cls_counts = np.bincount(y_train, minlength=COUNT_CLASSES).astype(np.float32)
+    cls_counts = np.where(cls_counts > 0, cls_counts, 1.0)
+    cls_weight = (1.0 / cls_counts) / (1.0 / cls_counts).sum() * COUNT_CLASSES
+    cls_weight_t = torch.from_numpy(cls_weight).to(device)
+    print(f"class weights: {cls_weight.tolist()}")
+
+    Xt = torch.from_numpy(X_train).to(device)
+    yt = torch.from_numpy(y_train).to(device)
+    Xe = torch.from_numpy(X_eval).to(device)
+    ye = torch.from_numpy(y_eval).to(device)
+
+    model = CountNet().to(device)
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=50, T_mult=1)
+
+    n_train = X_train.shape[0]
+    epoch_losses = []
+    t0 = time.perf_counter()
+
+    best_eval_acc = 0.0
+    best_state = None
+
+    for epoch in range(args.epochs):
+        model.train()
+        perm = torch.randperm(n_train, device=device)
+        train_loss = 0.0
+        train_correct = 0
+        n_batches = 0
+        for i in range(0, n_train, args.batch_size):
+            idx = perm[i : i + args.batch_size]
+            xb = Xt[idx]
+            yb = yt[idx]
+            opt.zero_grad()
+            count_logits, conf_logits = model(xb)
+
+            # Categorical cross-entropy for count.
+            ce = F.cross_entropy(count_logits, yb, weight=cls_weight_t)
+
+            # Confidence head: train against `argmax == truth` indicator.
+            with torch.no_grad():
+                pred = count_logits.argmax(dim=1)
+                correct_indicator = (pred == yb).float().unsqueeze(1)
+            bce = F.binary_cross_entropy_with_logits(conf_logits, correct_indicator)
+
+            # Brier-score uncertainty calibration on the conf head — sharpens
+            # the calibration so the sigmoid output is a real probability.
+            with torch.no_grad():
+                conf_sigm = torch.sigmoid(conf_logits)
+            brier = ((conf_sigm - correct_indicator) ** 2).mean()
+
+            loss = ce + 0.3 * bce + 0.1 * brier
+            loss.backward()
+            opt.step()
+
+            train_loss += loss.item()
+            train_correct += (pred == yb).sum().item()
+            n_batches += 1
+
+        sched.step()
+
+        model.eval()
+        with torch.no_grad():
+            cl_e, _ = model(Xe)
+            eval_loss = F.cross_entropy(cl_e, ye, weight=cls_weight_t).item()
+            eval_pred = cl_e.argmax(dim=1)
+            eval_acc = (eval_pred == ye).float().mean().item()
+            eval_within1 = ((eval_pred - ye).abs() <= 1).float().mean().item()
+
+        epoch_losses.append({
+            "epoch": epoch,
+            "train_loss": train_loss / n_batches,
+            "train_acc": train_correct / n_train,
+            "eval_loss": eval_loss,
+            "eval_acc": eval_acc,
+            "eval_within_pm1": eval_within1,
+        })
+
+        if eval_acc > best_eval_acc:
+            best_eval_acc = eval_acc
+            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
+
+        if epoch < 5 or epoch % 50 == 0 or epoch == args.epochs - 1:
+            print(f"epoch {epoch:3d}  train_loss={train_loss/n_batches:.4f}  "
+                  f"train_acc={train_correct/n_train:.3f}  "
+                  f"eval_loss={eval_loss:.4f}  eval_acc={eval_acc:.3f}  "
+                  f"within±1={eval_within1:.3f}")
+
+    train_time = time.perf_counter() - t0
+    print(f"\ntrained {args.epochs} epochs in {train_time:.1f} s")
+    print(f"best eval_acc: {best_eval_acc:.3f}")
+
+    # Restore best checkpoint
+    if best_state is not None:
+        model.load_state_dict(best_state)
+
+    # Eval breakdown
+    model.eval()
+    with torch.no_grad():
+        cl_e, conf_e = model(Xe)
+        probs_e = torch.softmax(cl_e, dim=1)
+        pred_e = cl_e.argmax(dim=1)
+        acc = (pred_e == ye).float().mean().item()
+        within1 = ((pred_e - ye).abs() <= 1).float().mean().item()
+        mae = (pred_e - ye).abs().float().mean().item()
+
+        # Per-class accuracy
+        per_class = {}
+        for k in range(COUNT_CLASSES):
+            mask = ye == k
+            n = mask.sum().item()
+            if n > 0:
+                per_class[k] = {
+                    "support": int(n),
+                    "accuracy": ((pred_e == ye) & mask).sum().item() / n,
+                }
+
+        # Confidence-accuracy calibration: Spearman over (predicted-correct, confidence)
+        conf_sigm = torch.sigmoid(conf_e).squeeze(-1)
+        correct = (pred_e == ye).float()
+        # Spearman = Pearson over ranks
+        c_rank = conf_sigm.argsort().argsort().float()
+        r_rank = correct.argsort().argsort().float()
+        c_centered = c_rank - c_rank.mean()
+        r_centered = r_rank - r_rank.mean()
+        denom = (c_centered.norm() * r_centered.norm()).item()
+        spearman = (c_centered * r_centered).sum().item() / denom if denom > 0 else 0.0
+
+    print(f"\n=== final eval ===")
+    print(f"  accuracy:       {acc:.3f}")
+    print(f"  within ±1:      {within1:.3f}")
+    print(f"  MAE:            {mae:.3f}")
+    print(f"  conf↔correct Spearman: {spearman:.3f}")
+    for k, v in per_class.items():
+        print(f"  class {k}:  {v['accuracy']:.3f} accuracy on {v['support']} samples")
+
+    # Save safetensors
+    write_safetensors(model, Path(args.out_safetensors))
+    print(f"\nwrote {args.out_safetensors} ({Path(args.out_safetensors).stat().st_size} bytes)")
+
+    # ONNX export
+    dummy = torch.zeros(1, N_SUB, N_FRAMES, device=device)
+    try:
+        torch.onnx.export(
+            model, dummy, args.out_onnx,
+            opset_version=18,
+            input_names=["csi_window"],
+            output_names=["count_logits", "conf_logits"],
+            dynamic_axes={
+                "csi_window": {0: "batch"},
+                "count_logits": {0: "batch"},
+                "conf_logits": {0: "batch"},
+            },
+            export_params=True,
+            do_constant_folding=True,
+        )
+        print(f"wrote {args.out_onnx} ({Path(args.out_onnx).stat().st_size} bytes)")
+    except Exception as e:
+        print(f"WARN: ONNX export failed: {e}")
+
+    # Results JSON
+    results = {
+        "backend": "candle-cuda" if device.type == "cuda" else "candle-cpu",
+        "device": str(device),
+        "epochs": args.epochs,
+        "train_time_s": train_time,
+        "best_eval_acc": best_eval_acc,
+        "final_eval_acc": acc,
+        "final_eval_within_pm1": within1,
+        "final_eval_mae": mae,
+        "conf_correctness_spearman": spearman,
+        "per_class_accuracy": per_class,
+        "hyperparameters": {
+            "optimizer": "AdamW",
+            "lr": args.lr,
+            "weight_decay": args.weight_decay,
+            "batch_size": args.batch_size,
+            "schedule": "cosine_warm_restarts",
+            "epochs": args.epochs,
+            "loss": "cross_entropy(count) + 0.3*bce(conf) + 0.1*brier(conf)",
+            "z_score_normalisation": True,
+            "class_weights": cls_weight.tolist(),
+        },
+        "epoch_losses": epoch_losses,
+    }
+    Path(args.out_results).write_text(json.dumps(results, indent=2))
+    print(f"wrote {args.out_results} ({Path(args.out_results).stat().st_size} bytes)")
+
+
+if __name__ == "__main__":
+    main()