ruvnet--RuView/benchmarks/wiflow-std/quantize_bench.py

"""ADR-152 "optimize beyond SOTA": edge-optimization benchmark for the
retrained WiFlow-STD checkpoint (results/retrained_best_pose_model.pth,
~96% PCK@20, fp32 params 2,225,042).

Measures, for fp32 / fp16 / dynamic-int8 torch variants:
  (a) serialized state_dict size on disk,
  (b) CPU inference latency per window at batch 1 and batch 64
      (median of repeated runs, this Windows box),
  (c) accuracy (PCK@20/50 + MPJPE, upstream metrics) on a corruption-free
      random subset of the seed-42 file-level 70/15/15 test split
      (same split as eval_repro.py; corrupted windows 487-499 excluded via
      results/nan_windows_mask.npy | results/big_windows_mask.npy).

Also verifies the paper's "~2.2 MB int8" size claim: reports which layer
types torch dynamic quantization actually converts (the model contains NO
nn.Linear -- it is Conv1d/Conv2d/BatchNorm only) and the real on-disk size.

Usage:
  .venv/Scripts/python.exe quantize_bench.py \
      --data-dir C:/Users/ruv/.cache/kagglehub/datasets/kaka2434/wiflow-dataset/versions/1/preprocessed_csi_data \
      [--subset 10000] [--skip-accuracy]

Writes/merges into results/edge_optimization.json under key "torch".
"""

import argparse
import json
import os
import platform
import statistics
import sys
import time

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

HERE = os.path.dirname(os.path.abspath(__file__))
UPSTREAM = os.path.join(HERE, "upstream")
RESULTS = os.path.join(HERE, "results")
sys.path.insert(0, UPSTREAM)

# Upstream models/__init__.py is broken as published (imports a name tcn.py
# does not define); register a stub package so it never executes.
import types  # noqa: E402

_models_pkg = types.ModuleType("models")
_models_pkg.__path__ = [os.path.join(UPSTREAM, "models")]
sys.modules["models"] = _models_pkg

import dataset as upstream_dataset  # noqa: E402
from dataset import (  # noqa: E402
    PreprocessedCSIKeypointsDataset,
    create_preprocessed_train_val_test_loaders,
)
from models.pose_model import WiFlowPoseModel  # noqa: E402
from utils.metrics import calculate_mpjpe, calculate_pck  # noqa: E402

CHECKPOINT = os.path.join(RESULTS, "retrained_best_pose_model.pth")

# csi_windows.npy is ~13 GB; mmap large arrays instead of loading into RAM
# (same trick as eval_repro.py).
_np_load = np.load


def _np_load_mmap(path, *a, **kw):
    if (isinstance(path, str) and path.endswith(".npy")
            and os.path.getsize(path) > 1 << 30 and "mmap_mode" not in kw):
        kw["mmap_mode"] = "r"
    return _np_load(path, *a, **kw)


upstream_dataset.np.load = _np_load_mmap


def load_fp32_model():
    state = torch.load(CHECKPOINT, map_location="cpu", weights_only=True)
    # legacy upstream names, harmless no-op on the retrained checkpoint
    renames = {"att.": "attention.", "final_conv.": "decoder."}
    state = {next((new + k[len(old):] for old, new in renames.items()
                   if k.startswith(old)), k): v
             for k, v in state.items()}
    model = WiFlowPoseModel(dropout=0.5)
    model.load_state_dict(state, strict=True)
    model.eval()
    return model


def state_dict_size_bytes(model, path):
    torch.save(model.state_dict(), path)
    return os.path.getsize(path)


def bench_latency(model, batch_size, n_runs, dtype=torch.float32):
    gen = torch.Generator().manual_seed(123)
    x = torch.rand(batch_size, 540, 20, generator=gen).to(dtype)
    with torch.no_grad():
        for _ in range(max(5, n_runs // 10)):  # warmup
            model(x)
        times = []
        for _ in range(n_runs):
            t0 = time.perf_counter()
            model(x)
            times.append(time.perf_counter() - t0)
    med = statistics.median(times)
    return {
        "batch_size": batch_size,
        "runs": n_runs,
        "median_ms_per_batch": med * 1e3,
        "median_ms_per_window": med * 1e3 / batch_size,
        "windows_per_second": batch_size / med,
    }


def build_test_subset(data_dir, subset_size, batch_size=64):
    """Seed-42 file-level 70/15/15 test split (exactly as eval_repro.py),
    minus corrupted windows, then a seed-42 random subset."""
    dataset = PreprocessedCSIKeypointsDataset(
        data_dir=data_dir, keypoint_scale=1000.0, enable_temporal_clean=True)
    _tr, _va, test_loader = create_preprocessed_train_val_test_loaders(
        dataset=dataset, batch_size=batch_size, num_workers=0, random_seed=42)
    test_indices = np.asarray(test_loader.dataset.indices)

    corrupted = (np.load(os.path.join(RESULTS, "nan_windows_mask.npy"))
                 | np.load(os.path.join(RESULTS, "big_windows_mask.npy")))
    clean = test_indices[~corrupted[test_indices]]
    print(f"test split: {len(test_indices)} windows, "
          f"{len(test_indices) - len(clean)} corrupted excluded, "
          f"{len(clean)} clean")

    if subset_size and subset_size < len(clean):
        rng = np.random.default_rng(42)
        clean = np.sort(rng.choice(clean, size=subset_size, replace=False))
    subset = torch.utils.data.Subset(dataset, clean.tolist())
    loader = DataLoader(subset, batch_size=batch_size, shuffle=False,
                        num_workers=0)
    return loader, len(clean)


def evaluate(model, loader, dtype=torch.float32, label=""):
    totals = {0.2: 0.0, 0.5: 0.0}
    total_mpe, n = 0.0, 0
    t0 = time.time()
    with torch.no_grad():
        for batch_idx, (bx, by) in enumerate(loader):
            out = model(bx.to(dtype)).float()
            pck = calculate_pck(out, by, thresholds=[0.2, 0.5])
            mpe = calculate_mpjpe(out, by)
            bs = by.size(0)
            total_mpe += mpe * bs
            for t in totals:
                totals[t] += pck[t] * bs
            n += bs
            if batch_idx % 50 == 0:
                print(f"  [{label}] batch {batch_idx}: n={n} "
                      f"pck20={totals[0.2]/n:.4f} mpjpe={total_mpe/n:.4f} "
                      f"({time.time()-t0:.0f}s)", flush=True)
    return {
        "samples": n,
        "pck@20": totals[0.2] / n,
        "pck@50": totals[0.5] / n,
        "mpjpe": total_mpe / n,
        "wall_seconds": time.time() - t0,
    }


def quantize_int8_dynamic(fp32_model):
    """torch.ao.quantization.quantize_dynamic on Linear/Conv where supported.
    Returns (model, report) where report documents what actually quantized."""
    qmodel = torch.ao.quantization.quantize_dynamic(
        fp32_model, {nn.Linear, nn.Conv1d, nn.Conv2d}, dtype=torch.qint8)

    quantized, total_params, quant_params = [], 0, 0
    for name, mod in qmodel.named_modules():
        cls = type(mod).__module__ + "." + type(mod).__name__
        if "quantized" in cls:
            w = mod.weight() if callable(getattr(mod, "weight", None)) else None
            numel = w.numel() if w is not None else 0
            quant_params += numel
            quantized.append({"module": name, "class": cls, "params": numel})
    for p in fp32_model.parameters():
        total_params += p.numel()

    n_linear = sum(isinstance(m, nn.Linear) for m in fp32_model.modules())
    n_conv1d = sum(isinstance(m, nn.Conv1d) for m in fp32_model.modules())
    n_conv2d = sum(isinstance(m, nn.Conv2d) for m in fp32_model.modules())
    report = {
        "eligible_module_counts": {
            "nn.Linear": n_linear, "nn.Conv1d": n_conv1d, "nn.Conv2d": n_conv2d},
        "modules_actually_quantized": quantized,
        "n_modules_quantized": len(quantized),
        "params_total": total_params,
        "params_quantized": quant_params,
        "params_quantized_fraction": quant_params / total_params,
    }
    return qmodel, report


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", default=os.path.join(
        os.path.expanduser("~"), ".cache", "kagglehub", "datasets", "kaka2434",
        "wiflow-dataset", "versions", "1", "preprocessed_csi_data"))
    parser.add_argument("--subset", type=int, default=10000)
    parser.add_argument("--runs-b1", type=int, default=100)
    parser.add_argument("--runs-b64", type=int, default=30)
    parser.add_argument("--skip-accuracy", action="store_true")
    parser.add_argument("--out", default=os.path.join(RESULTS, "edge_optimization.json"))
    args = parser.parse_args()

    torch.manual_seed(42)
    results = {
        "env": {
            "torch": torch.__version__,
            "platform": platform.platform(),
            "processor": platform.processor(),
            "num_threads": torch.get_num_threads(),
            "checkpoint": os.path.relpath(CHECKPOINT, HERE),
        },
        "variants": {},
    }

    # ---- build variants ---------------------------------------------------
    fp32 = load_fp32_model()
    n_params = sum(p.numel() for p in fp32.parameters())
    results["env"]["params"] = n_params
    print(f"fp32 model: {n_params:,} params")

    fp16 = load_fp32_model().half()

    int8, q_report = quantize_int8_dynamic(load_fp32_model())
    results["int8_dynamic_quant_report"] = q_report
    print(f"int8 dynamic: {q_report['n_modules_quantized']} modules quantized, "
          f"{q_report['params_quantized_fraction']*100:.1f}% of params")

    variants = {
        "fp32": (fp32, torch.float32, "retrained_fp32_resaved.pth"),
        "fp16": (fp16, torch.float16, "retrained_fp16.pth"),
        "int8_dynamic": (int8, torch.float32, "retrained_int8_dynamic.pth"),
    }

    # ---- (a) size + (b) latency -------------------------------------------
    for name, (model, dtype, fname) in variants.items():
        path = os.path.join(RESULTS, fname)
        size = state_dict_size_bytes(model, path)
        print(f"\n=== {name}: {size/1e6:.3f} MB on disk ({fname}) ===")
        lat1 = bench_latency(model, 1, args.runs_b1, dtype)
        lat64 = bench_latency(model, 64, args.runs_b64, dtype)
        print(f"  batch 1:  {lat1['median_ms_per_window']:.2f} ms/window "
              f"({lat1['windows_per_second']:.0f}/s)")
        print(f"  batch 64: {lat64['median_ms_per_window']:.3f} ms/window "
              f"({lat64['windows_per_second']:.0f}/s)")
        results["variants"][name] = {
            "file": fname,
            "size_bytes": size,
            "size_mb": size / 1e6,
            "latency_batch1": lat1,
            "latency_batch64": lat64,
        }

    # ---- (c) accuracy ------------------------------------------------------
    if not args.skip_accuracy:
        loader, n_clean = build_test_subset(args.data_dir, args.subset)
        results["accuracy_subset"] = {
            "description": "seed-42 file-level 70/15/15 test split, corrupted "
                           "windows (files 487-499) excluded, seed-42 random "
                           "subset",
            "subset_size": min(args.subset, n_clean) if args.subset else n_clean,
            "clean_test_total": n_clean,
        }
        for name, (model, dtype, _f) in variants.items():
            print(f"\n=== accuracy: {name} ===")
            results["variants"][name]["accuracy"] = evaluate(
                model, loader, dtype, label=name)
            print(json.dumps(results["variants"][name]["accuracy"], indent=2))

    # ---- merge into edge_optimization.json ---------------------------------
    merged = {}
    if os.path.exists(args.out):
        with open(args.out) as f:
            merged = json.load(f)
    merged["torch"] = results
    with open(args.out, "w") as f:
        json.dump(merged, f, indent=2)
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()