#!/usr/bin/env python3 """ Device-Category 1D-CNN — Phase 3B (PLAN_TO_PROD) ================================================ Trains a 1D convolutional net that predicts a device *category* directly from the normalized RAW pulse array of a Sub-GHz capture — the "shape of the signal" rather than the hand-picked timing statistics the Phase 3A model uses. Architecture (per PLAN_TO_PROD §Phase B): Conv1D(64,3) → BN → ReLU Conv1D(128,3) → BN → ReLU Conv1D(256,3) → BN → ReLU GlobalAvgPool → Dense(256) → Dropout(0.3) → Dense(N) → Softmax Input: pulse array scaled to [-1,1], padded/truncated to 512 (pulse_encoder). Honesty guardrails (identical discipline to the Phase 3A trainer): * GROUP-AWARE split — a device sub-folder never spans train/test, so we don't score inflated accuracy off near-duplicate repeat captures. * Epoch selection uses a validation slice carved from TRAIN only; the group-disjoint TEST set is scored exactly once at the end. * The heuristic router AND the Phase 3A statistical model are scored on the *same* held-out files → a fair three-way comparison, not cherry-picked. * Class-weighted loss + balanced accuracy reported, because the RAW corpus is ~70% Garage/Gate and raw accuracy alone is misleading. Run: python scripts/train_cnn_classifier.py python scripts/train_cnn_classifier.py --per-category 400 --epochs 60 """ import argparse import json import sys import time from collections import Counter from pathlib import Path import numpy as np sys.path.insert(0, str(Path(__file__).parent.parent)) from src.parser.sub_parser import parse_sub_file from src.matcher.timing_analyzer import get_timing_analyzer from src.matcher.preamble_detector import get_preamble_detector from src.matcher.category_router import get_category_router from src.matcher.pulse_encoder import encode_pulses, PULSE_LEN from scripts.benchmark_realworld import FOLDER_TO_CATEGORY, DEFAULT_DATASET from scripts.train_category_classifier import ( extract_features, heuristic_top1, FEATURE_NAMES, ) def collect(dataset_root: Path, per_category: int, seed: int): """Return aligned CNN arrays + stat features + labels/groups/freqs. Every accepted RAW file yields BOTH the CNN encoding and the Phase-3A 14-feature vector, indexed identically, so the same group split scores all three models on the same files. """ import random rng = random.Random(seed) ta, pd = get_timing_analyzer(), get_preamble_detector() Xcnn, Xstat, y, groups, freqs = [], [], [], [], [] skipped = Counter() for folder, category in FOLDER_TO_CATEGORY.items(): folder_path = dataset_root / folder if not folder_path.is_dir(): continue subs = sorted(folder_path.rglob("*.sub")) rng.shuffle(subs) taken = 0 for p in subs: if taken >= per_category: break try: meta = parse_sub_file(str(p)) except Exception: skipped["parse_error"] += 1 continue if not getattr(meta, "has_raw_data", False): skipped["no_raw"] += 1 continue enc = encode_pulses(meta.raw_data) if enc is None: skipped["bad_pulses"] += 1 continue feats = extract_features(meta.raw_data, meta.frequency, ta, pd) if feats is None: skipped["no_timing"] += 1 continue Xcnn.append(enc) Xstat.append(feats) y.append(category) groups.append(str(p.parent)) freqs.append(meta.frequency) taken += 1 return (np.asarray(Xcnn, dtype=np.float32), np.asarray(Xstat, dtype=np.float32), np.array(y), np.array(groups), np.array(freqs), skipped) def build_model(n_classes: int): import torch.nn as nn class PulseCNN(nn.Module): def __init__(self, n_out): super().__init__() self.features = nn.Sequential( nn.Conv1d(1, 64, 3, padding=1), nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(2), nn.Conv1d(64, 128, 3, padding=1), nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2), nn.Conv1d(128, 256, 3, padding=1), nn.BatchNorm1d(256), nn.ReLU(), nn.AdaptiveAvgPool1d(1), # GlobalAvgPool ) self.head = nn.Sequential( nn.Flatten(), nn.Linear(256, 256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, n_out), ) def forward(self, x): return self.head(self.features(x)) return PulseCNN(n_classes) def main(): ap = argparse.ArgumentParser() ap.add_argument("--dataset", type=Path, default=DEFAULT_DATASET) ap.add_argument("--per-category", type=int, default=400) ap.add_argument("--seed", type=int, default=42) ap.add_argument("--test-size", type=float, default=0.25) ap.add_argument("--val-size", type=float, default=0.15) ap.add_argument("--epochs", type=int, default=60) ap.add_argument("--batch-size", type=int, default=64) ap.add_argument("--lr", type=float, default=1e-3) ap.add_argument("--out", type=Path, default=Path("models")) args = ap.parse_args() if not args.dataset.is_dir(): print(f"Dataset not found: {args.dataset}", file=sys.stderr) sys.exit(1) import torch import torch.nn as nn from sklearn.model_selection import GroupShuffleSplit from sklearn.metrics import (accuracy_score, balanced_accuracy_score, f1_score, confusion_matrix, classification_report) torch.manual_seed(args.seed) np.random.seed(args.seed) print("Collecting + encoding RAW pulses ...") t0 = time.time() Xcnn, Xstat, y, groups, freqs, skipped = collect( args.dataset, args.per_category, args.seed) print(f" {len(Xcnn)} RAW samples in {time.time()-t0:.1f}s (skipped: {dict(skipped)})") print(f" class balance: {dict(Counter(y))}") print(f" distinct device groups: {len(set(groups))}") classes = sorted(set(y)) cls_to_i = {c: i for i, c in enumerate(classes)} yi = np.array([cls_to_i[c] for c in y]) # Group-aware test split (no device leakage) gss = GroupShuffleSplit(n_splits=1, test_size=args.test_size, random_state=args.seed) trainval_idx, test_idx = next(gss.split(Xcnn, yi, groups)) # Validation carved from train (for epoch selection only; test untouched) gss2 = GroupShuffleSplit(n_splits=1, test_size=args.val_size, random_state=args.seed) tr_rel, val_rel = next(gss2.split(Xcnn[trainval_idx], yi[trainval_idx], groups[trainval_idx])) train_idx = trainval_idx[tr_rel] val_idx = trainval_idx[val_rel] print(f" train={len(train_idx)} val={len(val_idx)} test={len(test_idx)} " f"(test groups={len(set(groups[test_idx]))}, group-disjoint)") def to_tensor(idx): x = torch.from_numpy(Xcnn[idx]).unsqueeze(1) # (N,1,L) t = torch.from_numpy(yi[idx]).long() return x, t Xtr, ytr = to_tensor(train_idx) Xval, yval = to_tensor(val_idx) Xte, yte = to_tensor(test_idx) # Class weights (inverse freq) for imbalance counts = np.bincount(yi[train_idx], minlength=len(classes)).astype(float) counts[counts == 0] = 1.0 weights = torch.tensor((counts.sum() / counts), dtype=torch.float32) weights = weights / weights.mean() model = build_model(len(classes)) opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) lossf = nn.CrossEntropyLoss(weight=weights) n = len(train_idx) best_val, best_state = -1.0, None for epoch in range(args.epochs): model.train() perm = torch.randperm(n) for i in range(0, n, args.batch_size): bi = perm[i:i + args.batch_size] opt.zero_grad() out = model(Xtr[bi]) loss = lossf(out, ytr[bi]) loss.backward() opt.step() # Validation balanced accuracy model.eval() with torch.no_grad(): vpred = model(Xval).argmax(1).numpy() vbal = balanced_accuracy_score(yval.numpy(), vpred) if vbal > best_val: best_val = vbal best_state = {k: v.clone() for k, v in model.state_dict().items()} if (epoch + 1) % 10 == 0: print(f" epoch {epoch+1:3d} val_bal_acc={vbal:.1%} best={best_val:.1%}") model.load_state_dict(best_state) model.eval() with torch.no_grad(): te_logits = model(Xte) te_pred = te_logits.argmax(1).numpy() yte_np = yte.numpy() acc = accuracy_score(yte_np, te_pred) bal = balanced_accuracy_score(yte_np, te_pred) f1 = f1_score(yte_np, te_pred, average="macro") # ── Baselines on the SAME test files ─────────────────────────────────── router = get_category_router() h_top1 = 0 for row, freq, truth in zip(Xstat[test_idx], freqs[test_idx], y[test_idx]): cat, _allowed, _full = heuristic_top1(row, freq, router) if cat == truth: h_top1 += 1 h_acc = h_top1 / len(test_idx) stat_acc = stat_bal = None try: import joblib bundle = joblib.load(args.out / "category_classifier.joblib") stat_model, stat_classes = bundle["model"], [str(c) for c in bundle["classes"]] stat_pred = stat_model.predict(Xstat[test_idx]) stat_acc = accuracy_score(y[test_idx], stat_pred) stat_bal = balanced_accuracy_score(y[test_idx], stat_pred) except Exception as e: print(f" (statistical baseline unavailable: {e})") print("\n" + "=" * 66) print("THREE-WAY COMPARISON (same group-disjoint test set)") print("=" * 66) print(f"{'model':26}{'top-1':>10}{'balanced':>11}") print(f"{'heuristic router':26}{h_acc:>10.1%}{'—':>11}") if stat_acc is not None: print(f"{'statistical (RF/GB)':26}{stat_acc:>10.1%}{stat_bal:>11.1%}") print(f"{'1D CNN (this run)':26}{acc:>10.1%}{bal:>11.1%}") print(f"\nCNN macro F1: {f1:.1%}") print("\nPer-category (CNN):") print(classification_report(yte_np, te_pred, labels=list(range(len(classes))), target_names=classes, zero_division=0)) print("Confusion (rows=truth, cols=pred): labels=", classes) print(confusion_matrix(yte_np, te_pred, labels=list(range(len(classes))))) # ── Persist: torch state + ONNX + metrics ────────────────────────────── args.out.mkdir(parents=True, exist_ok=True) torch.save({"state_dict": best_state, "classes": classes, "pulse_len": PULSE_LEN}, args.out / "category_cnn.pt") dummy = torch.zeros(1, 1, PULSE_LEN) onnx_path = args.out / "category_cnn.onnx" try: torch.onnx.export( model, dummy, str(onnx_path), input_names=["pulses"], output_names=["logits"], dynamic_axes={"pulses": {0: "batch"}, "logits": {0: "batch"}}, opset_version=13, ) onnx_ok = True except Exception as e: onnx_ok = False print(f" (ONNX export failed, torch .pt still saved: {e})") meta = { "n_samples": int(len(Xcnn)), "n_train": int(len(train_idx)), "n_val": int(len(val_idx)), "n_test": int(len(test_idx)), "classes": classes, "pulse_len": PULSE_LEN, "epochs": args.epochs, "class_balance": {k: int(v) for k, v in Counter(y).items()}, "cnn": {"accuracy": acc, "balanced_accuracy": bal, "macro_f1": f1}, "heuristic_same_test": {"top1": h_acc}, "statistical_same_test": ( {"top1": stat_acc, "balanced_accuracy": stat_bal} if stat_acc is not None else None), } (args.out / "category_cnn_metrics.json").write_text(json.dumps(meta, indent=2)) print(f"\nSaved torch -> {args.out / 'category_cnn.pt'}") if onnx_ok: print(f"Saved ONNX -> {onnx_path}") print(f"Saved metrics-> {args.out / 'category_cnn_metrics.json'}") if __name__ == "__main__": main()