/
GreyCat
/
RuView
Обзор
Документация
Войти
/
GreyCat
/
RuView
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
Аналитика
Безопасность
v1740
scripts/train-count.py
761 строка
32 KB
rUv
feat(cog-person-count): v0.0.2 — K-fold + label-smoothing + temperature-calibrated (#699)
22 май 2026, 02:47
Не верифицирован
22 май 2026, 02:47
b3a5012
Код
Авторство
О чём код?
#!/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()