mirror of
https://github.com/ruvnet/RuView
synced 2026-07-16 16:23:18 +00:00
feat(train): compact WiFlow-STD presets in Rust + tiny edge artifact (ADR-152)
WiFlowStdConfig gains half()/quarter()/tiny() mirroring the overnight sweep exactly: TcnGroupsMode (Fixed/Gcd/Depthwise), input_pw_groups, derived stride schedule and decoder-mid (all default to upstream behavior; legacy serde JSON unaffected). Param formulas pin to trained ground truth first try: 843,834 / 338,600 / 56,290; default 2,225,042 pin and 1.192e-7 parity unchanged. 248 tests green. Tiny edge artifact (tiny_edge_bench.py): ONNX fp32 = 295 KB, 0.66 ms/win (~1,500/s CPU), 94.11% PCK@20 (matches sweep clean-test exactly; parity 1.49e-7). Static int8 is a bad trade at this scale (-1.43pt, +19% MPJPE, -16% size, slower) — recorded as negative result. Export note: width-16 breaks AdaptiveAvgPool((15,1)) TorchScript export; replaced by exact mean+matmul equivalent, proven by parity. Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruv
@@ -267,6 +267,64 @@ Findings:
|
||||
says capacity *hurts* cross-subject, so the compact end may generalize no
|
||||
worse, but that is a hypothesis, not a measurement.
|
||||
|
||||
### Compact-variant edge artifacts (MEASURED, 2026-06-11)
|
||||
|
||||
Edge pipeline for the **tiny** checkpoint (56,290 params), same machinery and
|
||||
protocol as the full-model edge rows above (this Windows box, torch
|
||||
2.12.0+cpu, onnxruntime 1.26.0; dynamic-batch opset-17 TorchScript export;
|
||||
static QDQ **Percentile(99.99) conv-only** int8 calibrated on **512**
|
||||
corruption-free TRAIN-split windows; accuracy on the identical 10k-window
|
||||
seed-42 clean test subset; latency = median ms/window over 3 interleaved
|
||||
reps, with the full-model fp32/int8 sessions interleaved as same-session
|
||||
references). Script: `tiny_edge_bench.py`; raw:
|
||||
`results/edge_optimization.json` (`tiny_variant`). Torch-vs-ORT parity on the
|
||||
stored fixture input: **max abs diff 1.5e-7 — PASS** (< 1e-4). The tiny fp32
|
||||
subset PCK@20 (94.11%) matches the full clean-test sweep figure (94.11%)
|
||||
exactly, so the subset remains representative.
|
||||
|
||||
Two forced deviations, both recorded in the JSON:
|
||||
|
||||
1. **Adaptive-pool export rewrite.** tiny's derived stride schedule
|
||||
`[2,1,1,1]` leaves feature width 16, and the TorchScript exporter rejects
|
||||
`AdaptiveAvgPool2d((15,1))` when 15 is not a factor of the input height
|
||||
(the full model never hit this — its width was exactly 15). Since the
|
||||
pool over a fixed-size map is a fixed linear operator, the export wrapper
|
||||
replaces it with `mean(-1)` (W axis, a factor) + a constant averaging
|
||||
matmul using PyTorch's exact bin rule; the parity check (vs the original
|
||||
torch model with the real pool) proves exactness.
|
||||
2. **Calibration count 512, not "~500"**: ORT 1.26's histogram collector
|
||||
`np.asarray()`'s the per-batch maxima, so the calibration count must be a
|
||||
multiple of the 64-window calibration batch or the ragged last batch
|
||||
crashes it (the earlier static-PTQ run dodged this by using exactly 512).
|
||||
|
||||
| Variant | Disk size | Batch 1 (ms/win) | Batch 64 (ms/win) | PCK@20 | PCK@50 | MPJPE |
|
||||
|---|---|---|---|---|---|---|
|
||||
| full ONNX fp32 (same-session ref) | 8.97 MB | 2.27 | 1.42 | 96.68% | 99.15% | 0.00936 |
|
||||
| full static QDQ Percentile conv-only (same-session ref) | 2.53 MB | 5.53 | 3.82 | 96.61% | 99.16% | 0.01031 |
|
||||
| **tiny ONNX fp32** | **0.295 MB** | **0.66** | **0.24** | **94.11%** | 99.37% | 0.01253 |
|
||||
| tiny static QDQ Percentile conv-only | 0.248 MB | 0.85 | 1.03 | 92.68% | 99.33% | 0.01491 |
|
||||
|
||||
(tiny torch `.pth` checkpoint for reference: 0.34 MB on disk; 56,290 fp32
|
||||
params ≈ 225 KB of weights.)
|
||||
|
||||
Findings:
|
||||
|
||||
- **The smallest deployable WiFlow-class model is the tiny ONNX fp32
|
||||
artifact: ~295 KB on disk, 0.66 ms/window batch-1 CPU (~1,500 windows/s),
|
||||
94.1% PCK@20** — 30× smaller and ~3.4× faster (in-session) than the full
|
||||
ONNX fp32 model for −2.6 pt PCK@20.
|
||||
- **int8 is a bad trade at this scale.** Static QDQ conv-only — the recipe
|
||||
that cost the full model only 0.07 pt — costs tiny **−1.43 pt** PCK@20
|
||||
(94.11 → 92.68%) and +19% MPJPE, saves only 47 KB (−16%; QDQ scales and
|
||||
the fp32 BN/attention glue are proportionally larger in a small graph),
|
||||
and is *slower* than tiny fp32 (0.85 vs 0.66 ms b1; 1.03 vs 0.24 ms b64 —
|
||||
QDQ kernel overhead dominates when the convs are this small). A 56k-param
|
||||
model has little redundancy left to absorb weight+activation rounding.
|
||||
- Deployment guidance, compact edition: ship tiny as **ONNX fp32** — at
|
||||
295 KB the int8 size saving solves no real constraint and costs accuracy
|
||||
and speed. If ~250 KB vs ~295 KB ever matters, weight-only quantization
|
||||
would be the thing to try next, not QDQ.
|
||||
|
||||
## Measurement (b): BLOCKED-ON-DATA (attempted 2026-06-10)
|
||||
|
||||
The fine-tune-on-ESP32 measurement stopped at dataset characterization, per the
|
||||
|
||||
@@ -626,5 +626,147 @@
|
||||
"description": "seed-42 file-level 70/15/15 test split, corrupted windows excluded, seed-42 random subset (same as quantize_bench/eval_ort_accuracy)",
|
||||
"subset_size": 10000
|
||||
}
|
||||
},
|
||||
"tiny_variant": {
|
||||
"env": {
|
||||
"torch": "2.12.0+cpu",
|
||||
"onnxruntime": "1.26.0",
|
||||
"platform": "Windows-11-10.0.26200-SP0",
|
||||
"num_threads": 16,
|
||||
"checkpoint": "results\\tiny_best.pth",
|
||||
"checkpoint_size_bytes": 340555,
|
||||
"params": 56290,
|
||||
"variant_config": {
|
||||
"tcn": [
|
||||
68,
|
||||
56,
|
||||
44,
|
||||
32
|
||||
],
|
||||
"conv": [
|
||||
2,
|
||||
4,
|
||||
8,
|
||||
16
|
||||
],
|
||||
"attn_groups": 2,
|
||||
"groups_mode": "depthwise",
|
||||
"input_pw_groups": 4
|
||||
}
|
||||
},
|
||||
"export": {
|
||||
"mode": "dynamic-batch",
|
||||
"exporter": "torchscript",
|
||||
"opset": 17,
|
||||
"file": "tiny_fp32_dynamic.onnx",
|
||||
"size_bytes": 295279,
|
||||
"size_mb": 0.295279,
|
||||
"verified_batches": [
|
||||
1,
|
||||
2,
|
||||
64
|
||||
],
|
||||
"note": "AdaptiveAvgPool2d((15,1)) replaced at export by an exact mean(-1) + constant averaging matmul (final_width 16 is not a multiple of 15, which the TorchScript exporter rejects); exactness proven by the parity check vs the original torch model"
|
||||
},
|
||||
"parity": {
|
||||
"fixture": "results/parity_fixture.npz input (batch 2, seed 42); reference output recomputed with the tiny torch model",
|
||||
"max_abs_diff_vs_torch": 1.4901161193847656e-07,
|
||||
"pass_lt_1e-4": true
|
||||
},
|
||||
"int8_static_percentile_conv": {
|
||||
"file": "tiny_int8_static_percentile_conv.onnx",
|
||||
"size_bytes": 248278,
|
||||
"size_mb": 0.248278,
|
||||
"calibration": {
|
||||
"method": "percentile",
|
||||
"percentile": 99.99,
|
||||
"windows": 512,
|
||||
"scope": "conv-only TRAIN-split corruption-free",
|
||||
"seconds": 1.5347836017608643
|
||||
},
|
||||
"per_channel": true,
|
||||
"activation_type": "QInt8",
|
||||
"weight_type": "QInt8",
|
||||
"max_abs_diff_vs_fp32_fixture": 0.018491357564926147
|
||||
},
|
||||
"latency": {
|
||||
"note": "3 interleaved repetitions per variant, median ms/window; full-model sessions are same-session references",
|
||||
"tiny_onnx_fp32": {
|
||||
"batch1_reps": [
|
||||
0.6312500008789357,
|
||||
0.6834500018157996,
|
||||
0.6595999984710943
|
||||
],
|
||||
"batch64_reps": [
|
||||
0.37747578119251557,
|
||||
0.24196640623586063,
|
||||
0.2314671875183194
|
||||
],
|
||||
"batch1_ms_per_window_median": 0.6595999984710943,
|
||||
"batch64_ms_per_window_median": 0.24196640623586063
|
||||
},
|
||||
"tiny_onnx_int8_static_percentile_conv": {
|
||||
"batch1_reps": [
|
||||
0.7988500001374632,
|
||||
0.9382499993080273,
|
||||
0.8451000030618161
|
||||
],
|
||||
"batch64_reps": [
|
||||
0.9211476562995813,
|
||||
1.3045390625165965,
|
||||
1.026230468767153
|
||||
],
|
||||
"batch1_ms_per_window_median": 0.8451000030618161,
|
||||
"batch64_ms_per_window_median": 1.026230468767153
|
||||
},
|
||||
"full_onnx_fp32_reference": {
|
||||
"batch1_reps": [
|
||||
2.267249998112675,
|
||||
2.80170000041835,
|
||||
2.132149998942623
|
||||
],
|
||||
"batch64_reps": [
|
||||
1.3050578124875756,
|
||||
1.4244992187855132,
|
||||
1.8014164062947202
|
||||
],
|
||||
"batch1_ms_per_window_median": 2.267249998112675,
|
||||
"batch64_ms_per_window_median": 1.4244992187855132
|
||||
},
|
||||
"full_onnx_int8_static_percentile_conv_reference": {
|
||||
"batch1_reps": [
|
||||
5.529599999135826,
|
||||
4.768399998283712,
|
||||
6.215800000063609
|
||||
],
|
||||
"batch64_reps": [
|
||||
3.815724218725336,
|
||||
3.1025562500417436,
|
||||
4.333318749957016
|
||||
],
|
||||
"batch1_ms_per_window_median": 5.529599999135826,
|
||||
"batch64_ms_per_window_median": 3.815724218725336
|
||||
}
|
||||
},
|
||||
"accuracy_subset": {
|
||||
"description": "seed-42 file-level 70/15/15 test split, corrupted windows excluded, seed-42 random subset (same as quantize_bench/eval_ort_accuracy/static_ptq_bench)",
|
||||
"subset_size": 10000
|
||||
},
|
||||
"accuracy": {
|
||||
"tiny_onnx_fp32": {
|
||||
"samples": 10000,
|
||||
"pck@20": 0.941106667804718,
|
||||
"pck@50": 0.99369333152771,
|
||||
"mpjpe": 0.012527281279861927,
|
||||
"wall_seconds": 10.927234888076782
|
||||
},
|
||||
"tiny_onnx_int8_static_percentile_conv": {
|
||||
"samples": 10000,
|
||||
"pck@20": 0.9268133331298828,
|
||||
"pck@50": 0.9932933319091797,
|
||||
"mpjpe": 0.014906252065300942,
|
||||
"wall_seconds": 12.320892333984375
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
"""ADR-152 efficiency-sweep follow-up: edge pipeline for the TINY compact
|
||||
WiFlow-STD variant (56,290 params, results/tiny_best.pth, trained overnight
|
||||
2026-06-10/11 -- see RESULTS.md "Efficiency sweep").
|
||||
|
||||
Headline question: what does the smallest deployable WiFlow-class model look
|
||||
like (KB + ms + PCK)? Reuses the onnx_bench.py / static_ptq_bench.py
|
||||
machinery on the tiny checkpoint:
|
||||
|
||||
1. Load tiny_best.pth with remote/sweep/model_compact.py
|
||||
(depthwise TCN groups, input_pw_groups=4, conv [2,4,8,16], attn groups 2).
|
||||
2. Export ONNX: dynamic batch, opset 17, TorchScript exporter (dynamo=False)
|
||||
-- same recipe that worked for the full model; verified at batch 1/2/64.
|
||||
One forced deviation: tiny's stride schedule [2,1,1,1] leaves final_width
|
||||
16, and the TorchScript exporter cannot export AdaptiveAvgPool2d((15,1))
|
||||
when 15 is not a factor of the input height (the full model never hit
|
||||
this -- its width was exactly 15). The adaptive pool over a fixed-size
|
||||
feature map is a fixed linear map, so the export wrapper replaces it with
|
||||
an exact matmul equivalent (PyTorch adaptive-pool bin semantics:
|
||||
bin i averages rows floor(i*H/K)..ceil((i+1)*H/K)); the W axis (20->1,
|
||||
a factor) becomes mean(-1). Exactness is proven by the parity check
|
||||
below, which compares against the ORIGINAL torch model with the real
|
||||
AdaptiveAvgPool2d.
|
||||
3. Torch-vs-ORT parity on the stored fixture input
|
||||
(results/parity_fixture.npz, batch 2, seed 42 -- same 540x20 input layout;
|
||||
reference output recomputed with the tiny torch model). PASS < 1e-4.
|
||||
4. Static QDQ conv-only int8 (quant_pre_process + quantize_static,
|
||||
per-channel QInt8 weights+activations, Percentile(99.99) calibration on
|
||||
512 corruption-free TRAIN-split windows -- the winning recipe and
|
||||
calibration count from static_ptq_bench.py. 512, not "about 500":
|
||||
ORT 1.26's histogram collector np.asarray()'s the per-batch maxima, so
|
||||
the calibration count must be a multiple of the batch size 64 or the
|
||||
ragged last batch crashes it).
|
||||
5. Disk size + CPU latency b1/b64 (3 interleaved reps, median ms/window)
|
||||
for tiny fp32 + tiny int8, with the full-model ONNX fp32 + static-int8
|
||||
sessions interleaved as same-session references.
|
||||
6. Accuracy (PCK@20/50 + MPJPE) on the identical 10k-window seed-42
|
||||
corruption-free test subset for tiny fp32 + tiny int8.
|
||||
|
||||
Usage:
|
||||
PYTHONUTF8=1 .venv/Scripts/python.exe tiny_edge_bench.py \
|
||||
[--data-dir <preprocessed_csi_data>] [--subset 10000] [--calib 500]
|
||||
|
||||
Writes/merges into results/edge_optimization.json under key "tiny_variant".
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
RESULTS = os.path.join(HERE, "results")
|
||||
sys.path.insert(0, HERE)
|
||||
sys.path.insert(0, os.path.join(HERE, "remote", "sweep"))
|
||||
|
||||
# quantize_bench sets up upstream imports + the np.load mmap patch
|
||||
from quantize_bench import build_test_subset # noqa: E402
|
||||
from eval_ort_accuracy import evaluate_ort # noqa: E402
|
||||
from static_ptq_bench import ( # noqa: E402
|
||||
build_calibration_windows,
|
||||
interleaved_latency,
|
||||
make_reader,
|
||||
ort_session,
|
||||
)
|
||||
from model_compact import CompactWiFlowPoseModel, describe # noqa: E402
|
||||
|
||||
TINY_CKPT = os.path.join(RESULTS, "tiny_best.pth")
|
||||
TINY_FP32_ONNX = os.path.join(RESULTS, "tiny_fp32_dynamic.onnx")
|
||||
TINY_PREPROC_ONNX = os.path.join(RESULTS, "tiny_fp32_preproc.onnx")
|
||||
TINY_INT8_ONNX = os.path.join(RESULTS, "tiny_int8_static_percentile_conv.onnx")
|
||||
FULL_FP32_ONNX = os.path.join(RESULTS, "retrained_fp32_dynamic.onnx")
|
||||
FULL_INT8_ONNX = os.path.join(RESULTS, "retrained_int8_static_percentile_conv.onnx")
|
||||
|
||||
# Exact tiny config from remote/sweep/run_sweep.py VARIANTS (measured 56,290
|
||||
# params, clean-test PCK@20 94.11% -- results/efficiency_sweep.jsonl).
|
||||
TINY = dict(tcn=[68, 56, 44, 32], conv=[2, 4, 8, 16], attn_groups=2,
|
||||
groups_mode="depthwise", input_pw_groups=4)
|
||||
|
||||
|
||||
def load_tiny_model():
|
||||
model = CompactWiFlowPoseModel(
|
||||
tcn_channels=TINY["tcn"], conv_channels=TINY["conv"],
|
||||
attn_groups=TINY["attn_groups"], groups_mode=TINY["groups_mode"],
|
||||
input_pw_groups=TINY["input_pw_groups"], dropout=0.5)
|
||||
state = torch.load(TINY_CKPT, map_location="cpu", weights_only=True)
|
||||
model.load_state_dict(state, strict=True)
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def adaptive_pool_matrix(h_in, h_out):
|
||||
"""Exact AdaptiveAvgPool1d as a (h_out, h_in) averaging matrix, using
|
||||
PyTorch's bin rule: bin i covers rows floor(i*h_in/h_out) ..
|
||||
ceil((i+1)*h_in/h_out)."""
|
||||
w = torch.zeros(h_out, h_in)
|
||||
for i in range(h_out):
|
||||
s = (i * h_in) // h_out
|
||||
e = -((-(i + 1) * h_in) // h_out) # ceil division
|
||||
w[i, s:e] = 1.0 / (e - s)
|
||||
return w
|
||||
|
||||
|
||||
class ExportWrapper(torch.nn.Module):
|
||||
"""CompactWiFlowPoseModel forward with the AdaptiveAvgPool2d((K,1))
|
||||
replaced by an exact fixed linear map (mean over the factor W axis, then
|
||||
a constant averaging matmul over the non-factor H axis) so the
|
||||
TorchScript ONNX exporter accepts it. Bit-equivalent up to float
|
||||
round-off; proven by the parity check against the original model."""
|
||||
|
||||
def __init__(self, m, num_keypoints=15):
|
||||
super().__init__()
|
||||
self.m = m
|
||||
self.register_buffer(
|
||||
"pool_w_t", adaptive_pool_matrix(m.final_width, num_keypoints).t())
|
||||
|
||||
def forward(self, x):
|
||||
m = self.m
|
||||
x = m.tcn(x)
|
||||
x = x.transpose(1, 2).unsqueeze(1)
|
||||
x = m.up(x)
|
||||
for block in m.residual_blocks:
|
||||
x = block(x)
|
||||
x = x.permute(0, 1, 3, 2)
|
||||
x = m.attention(x)
|
||||
x = m.decoder(x) # [B, 2, H=final_width, T=20]
|
||||
x = x.mean(-1) # W-axis pool (20 -> 1, a factor)
|
||||
x = x.matmul(self.pool_w_t) # exact adaptive H pool: [B, 2, K]
|
||||
return x.transpose(1, 2) # [B, K, 2]
|
||||
|
||||
|
||||
def export_onnx(model):
|
||||
"""Dynamic-batch TorchScript export (the recipe that worked for the full
|
||||
model in onnx_bench.py), verified at batch 1/2/64. Uses ExportWrapper
|
||||
(see docstring) because final_width 16 is not a multiple of 15."""
|
||||
wrapper = ExportWrapper(model).eval()
|
||||
x = torch.rand(2, 540, 20)
|
||||
with torch.no_grad():
|
||||
torch.onnx.export(
|
||||
wrapper, (x,), TINY_FP32_ONNX, opset_version=17,
|
||||
input_names=["input"], output_names=["output"], dynamo=False,
|
||||
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})
|
||||
sess = ort_session(TINY_FP32_ONNX)
|
||||
inp = sess.get_inputs()[0].name
|
||||
for b in (1, 2, 64):
|
||||
y = sess.run(None, {inp: np.zeros((b, 540, 20), dtype=np.float32)})[0]
|
||||
assert y.shape == (b, 15, 2), y.shape
|
||||
return {
|
||||
"mode": "dynamic-batch", "exporter": "torchscript", "opset": 17,
|
||||
"file": os.path.basename(TINY_FP32_ONNX),
|
||||
"size_bytes": os.path.getsize(TINY_FP32_ONNX),
|
||||
"size_mb": os.path.getsize(TINY_FP32_ONNX) / 1e6,
|
||||
"verified_batches": [1, 2, 64],
|
||||
"note": "AdaptiveAvgPool2d((15,1)) replaced at export by an exact "
|
||||
"mean(-1) + constant averaging matmul (final_width 16 is not "
|
||||
"a multiple of 15, which the TorchScript exporter rejects); "
|
||||
"exactness proven by the parity check vs the original torch "
|
||||
"model",
|
||||
}
|
||||
|
||||
|
||||
def quantize_tiny(calib_windows):
|
||||
"""quant_pre_process + static QDQ conv-only Percentile(99.99) int8 --
|
||||
the winning recipe from static_ptq_bench.py."""
|
||||
from onnxruntime.quantization import (CalibrationMethod, QuantFormat,
|
||||
QuantType, quantize_static)
|
||||
from onnxruntime.quantization.shape_inference import quant_pre_process
|
||||
|
||||
quant_pre_process(TINY_FP32_ONNX, TINY_PREPROC_ONNX)
|
||||
t0 = time.time()
|
||||
quantize_static(
|
||||
TINY_PREPROC_ONNX, TINY_INT8_ONNX, make_reader(calib_windows),
|
||||
quant_format=QuantFormat.QDQ,
|
||||
op_types_to_quantize=["Conv"],
|
||||
per_channel=True,
|
||||
activation_type=QuantType.QInt8,
|
||||
weight_type=QuantType.QInt8,
|
||||
calibrate_method=CalibrationMethod.Percentile,
|
||||
extra_options={"CalibPercentile": 99.99},
|
||||
)
|
||||
return {
|
||||
"file": os.path.basename(TINY_INT8_ONNX),
|
||||
"size_bytes": os.path.getsize(TINY_INT8_ONNX),
|
||||
"size_mb": os.path.getsize(TINY_INT8_ONNX) / 1e6,
|
||||
"calibration": {"method": "percentile", "percentile": 99.99,
|
||||
"windows": int(len(calib_windows)),
|
||||
"scope": "conv-only TRAIN-split corruption-free",
|
||||
"seconds": time.time() - t0},
|
||||
"per_channel": True,
|
||||
"activation_type": "QInt8",
|
||||
"weight_type": "QInt8",
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
import onnxruntime
|
||||
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("--calib", type=int, default=512,
|
||||
help="calibration windows; must be a multiple of the "
|
||||
"64-window calibration batch (ORT histogram "
|
||||
"collector rejects ragged batches)")
|
||||
parser.add_argument("--skip-accuracy", action="store_true")
|
||||
parser.add_argument("--out", default=os.path.join(RESULTS, "edge_optimization.json"))
|
||||
args = parser.parse_args()
|
||||
|
||||
model = load_tiny_model()
|
||||
info = describe(model)
|
||||
print(f"tiny model: {info['params']:,} params, tcn_groups={info['tcn_groups_per_block']}, "
|
||||
f"strides={info['conv_strides']}, final_width={info['final_width']}")
|
||||
assert info["params"] == 56290, info["params"]
|
||||
|
||||
results = {
|
||||
"env": {
|
||||
"torch": torch.__version__,
|
||||
"onnxruntime": onnxruntime.__version__,
|
||||
"platform": platform.platform(),
|
||||
"num_threads": torch.get_num_threads(),
|
||||
"checkpoint": os.path.relpath(TINY_CKPT, HERE),
|
||||
"checkpoint_size_bytes": os.path.getsize(TINY_CKPT),
|
||||
"params": info["params"],
|
||||
"variant_config": TINY,
|
||||
},
|
||||
}
|
||||
|
||||
# ---- export + parity ----------------------------------------------------
|
||||
print("\n=== ONNX export (dynamic batch, opset 17, torchscript) ===")
|
||||
results["export"] = export_onnx(model)
|
||||
print(f" {results['export']['size_mb']:.3f} MB, batches {results['export']['verified_batches']} OK")
|
||||
|
||||
fixture = np.load(os.path.join(RESULTS, "parity_fixture.npz"))
|
||||
fx = fixture["input"] # (2, 540, 20), seed 42 -- same input layout as full model
|
||||
sess_fp32 = ort_session(TINY_FP32_ONNX)
|
||||
y_ort = sess_fp32.run(None, {sess_fp32.get_inputs()[0].name: fx})[0]
|
||||
with torch.no_grad():
|
||||
y_torch = model(torch.from_numpy(fx)).numpy()
|
||||
results["parity"] = {
|
||||
"fixture": "results/parity_fixture.npz input (batch 2, seed 42); "
|
||||
"reference output recomputed with the tiny torch model",
|
||||
"max_abs_diff_vs_torch": float(np.abs(y_ort - y_torch).max()),
|
||||
"pass_lt_1e-4": bool(np.abs(y_ort - y_torch).max() < 1e-4),
|
||||
}
|
||||
print("parity:", json.dumps(results["parity"], indent=2))
|
||||
assert results["parity"]["pass_lt_1e-4"], "torch-vs-ORT parity FAILED"
|
||||
|
||||
# ---- static PTQ int8 ------------------------------------------------------
|
||||
print(f"\n=== static QDQ int8 (Percentile conv-only, {args.calib} calib windows) ===")
|
||||
calib = build_calibration_windows(args.data_dir, args.calib)
|
||||
results["int8_static_percentile_conv"] = quantize_tiny(calib)
|
||||
print(f" {results['int8_static_percentile_conv']['size_mb']:.3f} MB")
|
||||
sess_int8 = ort_session(TINY_INT8_ONNX)
|
||||
yq = sess_int8.run(None, {sess_int8.get_inputs()[0].name: fx})[0]
|
||||
results["int8_static_percentile_conv"]["max_abs_diff_vs_fp32_fixture"] = float(
|
||||
np.abs(yq - y_torch).max())
|
||||
|
||||
# ---- latency (3 interleaved reps, full-model sessions as references) -----
|
||||
print("\n=== latency (3 interleaved reps) ===")
|
||||
lat_sessions = {
|
||||
"tiny_onnx_fp32": sess_fp32,
|
||||
"tiny_onnx_int8_static_percentile_conv": sess_int8,
|
||||
"full_onnx_fp32_reference": ort_session(FULL_FP32_ONNX),
|
||||
"full_onnx_int8_static_percentile_conv_reference": ort_session(FULL_INT8_ONNX),
|
||||
}
|
||||
results["latency"] = {
|
||||
"note": "3 interleaved repetitions per variant, median ms/window; "
|
||||
"full-model sessions are same-session references",
|
||||
**interleaved_latency(lat_sessions),
|
||||
}
|
||||
|
||||
# ---- accuracy on the standard 10k corruption-free test subset ------------
|
||||
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 excluded, seed-42 random subset (same as "
|
||||
"quantize_bench/eval_ort_accuracy/static_ptq_bench)",
|
||||
"subset_size": min(args.subset, n_clean) if args.subset else n_clean,
|
||||
}
|
||||
results["accuracy"] = {}
|
||||
for name, sess in (("tiny_onnx_fp32", sess_fp32),
|
||||
("tiny_onnx_int8_static_percentile_conv", sess_int8)):
|
||||
print(f"\n=== accuracy: {name} ===")
|
||||
results["accuracy"][name] = evaluate_ort(sess, loader, name)
|
||||
print(json.dumps(results["accuracy"][name], 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["tiny_variant"] = results
|
||||
with open(args.out, "w") as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
print(f"\nwrote {args.out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -19,6 +19,43 @@ pub const TCN_KERNEL: usize = 3;
|
||||
/// forwarded to the TCN), so it is a constant here rather than a config field.
|
||||
pub const CONV_BLOCK_DROPOUT: f64 = 0.3;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TcnGroupsMode
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// How the group count of each depthwise-grouped TCN convolution is chosen
|
||||
/// (ADR-152 efficiency sweep, `benchmarks/wiflow-std/remote/sweep/model_compact.py`).
|
||||
///
|
||||
/// The upstream reference hardcodes `groups = 20`, which does not divide the
|
||||
/// compact variants' channel counts (e.g. 270, 135, 85). The sweep's rules:
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TcnGroupsMode {
|
||||
/// Every grouped conv uses [`WiFlowStdConfig::tcn_groups`] verbatim
|
||||
/// (upstream behavior; requires divisibility). Default.
|
||||
#[default]
|
||||
Fixed,
|
||||
/// Per-conv groups = `gcd(channels, tcn_groups)` — equals `tcn_groups`
|
||||
/// wherever the upstream choice is valid (incl. the 540-channel input
|
||||
/// conv) and falls back to the largest common divisor otherwise.
|
||||
/// The sweep's `gcd20` mode (`half` / `quarter` presets).
|
||||
Gcd,
|
||||
/// Per-conv groups = channels (fully depthwise; `tiny` preset).
|
||||
Depthwise,
|
||||
}
|
||||
|
||||
fn gcd(a: usize, b: usize) -> usize {
|
||||
let (mut a, mut b) = (a, b);
|
||||
while b != 0 {
|
||||
(a, b) = (b, a % b);
|
||||
}
|
||||
a
|
||||
}
|
||||
|
||||
fn default_input_pw_groups() -> usize {
|
||||
1
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// WiFlowStdConfig
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -47,9 +84,26 @@ pub struct WiFlowStdConfig {
|
||||
|
||||
/// Group count for the depthwise-grouped TCN convolutions. The reference
|
||||
/// hardcodes **20**; exposed so non-540 subcarrier layouts can keep the
|
||||
/// divisibility invariant. Default: **20**.
|
||||
/// divisibility invariant. Default: **20**. Interpreted per
|
||||
/// [`Self::tcn_groups_mode`]: the verbatim group count in `Fixed` mode,
|
||||
/// the gcd base in `Gcd` mode, ignored in `Depthwise` mode.
|
||||
pub tcn_groups: usize,
|
||||
|
||||
/// Group-selection rule for the TCN's grouped convolutions
|
||||
/// (ADR-152 efficiency sweep). Default: [`TcnGroupsMode::Fixed`]
|
||||
/// (upstream behavior — every grouped conv uses [`Self::tcn_groups`]).
|
||||
#[serde(default)]
|
||||
pub tcn_groups_mode: TcnGroupsMode,
|
||||
|
||||
/// Group count for the **first** TCN block's pointwise (1×1) and residual
|
||||
/// downsample convs (`subcarriers → tcn_channels[0]`). The sweep's `tiny`
|
||||
/// variant uses **4** to break the dense-540-input parameter floor
|
||||
/// (~117k params, which alone exceeds tiny's budget); every other config
|
||||
/// uses **1** (upstream behavior). Must divide both `subcarriers` and
|
||||
/// `tcn_channels[0]`. Default: **1**.
|
||||
#[serde(default = "default_input_pw_groups")]
|
||||
pub input_pw_groups: usize,
|
||||
|
||||
/// Output channels of the 2-D conv encoder blocks. The first entry is
|
||||
/// also `ConvBlock1`'s output; each subsequent block downsamples the
|
||||
/// subcarrier axis by 2. Default: **[8, 16, 32, 64]**.
|
||||
@@ -75,6 +129,8 @@ impl Default for WiFlowStdConfig {
|
||||
window: 20,
|
||||
tcn_channels: vec![540, 440, 340, 240],
|
||||
tcn_groups: 20,
|
||||
tcn_groups_mode: TcnGroupsMode::Fixed,
|
||||
input_pw_groups: 1,
|
||||
conv_channels: vec![8, 16, 32, 64],
|
||||
attention_groups: 8,
|
||||
keypoints: 15,
|
||||
@@ -93,6 +149,52 @@ impl WiFlowStdConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// **half** compact preset (ADR-152 efficiency sweep, trained
|
||||
/// 2026-06-10/11): **843,834** parameters (0.38×), clean-test PCK@20
|
||||
/// **96.62%** — strictly dominates the full reference on its own
|
||||
/// benchmark. Per-conv groups = `gcd(channels, 20)`; stride schedule
|
||||
/// derives to `[2, 2, 2, 1]`. See
|
||||
/// `benchmarks/wiflow-std/results/efficiency_sweep.jsonl`.
|
||||
pub fn half() -> Self {
|
||||
WiFlowStdConfig {
|
||||
tcn_channels: vec![270, 220, 170, 120],
|
||||
tcn_groups_mode: TcnGroupsMode::Gcd,
|
||||
conv_channels: vec![4, 8, 16, 32],
|
||||
attention_groups: 4,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// **quarter** compact preset (ADR-152 efficiency sweep): **338,600**
|
||||
/// parameters (0.15×), clean-test PCK@20 **96.05%**. Per-conv groups =
|
||||
/// `gcd(channels, 20)`; stride schedule derives to `[2, 2, 1, 1]`.
|
||||
pub fn quarter() -> Self {
|
||||
WiFlowStdConfig {
|
||||
tcn_channels: vec![135, 110, 85, 60],
|
||||
tcn_groups_mode: TcnGroupsMode::Gcd,
|
||||
conv_channels: vec![2, 4, 8, 16],
|
||||
attention_groups: 2,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// **tiny** compact preset (ADR-152 efficiency sweep): **56,290**
|
||||
/// parameters (0.025×), clean-test PCK@20 **94.11%** — the smallest
|
||||
/// deployable WiFlow-class model (~220 KB fp32). Fully depthwise TCN
|
||||
/// groups plus `input_pw_groups = 4` on the first block's pointwise /
|
||||
/// downsample convs; stride schedule derives to `[2, 1, 1, 1]`
|
||||
/// (feature width 16).
|
||||
pub fn tiny() -> Self {
|
||||
WiFlowStdConfig {
|
||||
tcn_channels: vec![68, 56, 44, 32],
|
||||
tcn_groups_mode: TcnGroupsMode::Depthwise,
|
||||
input_pw_groups: 4,
|
||||
conv_channels: vec![2, 4, 8, 16],
|
||||
attention_groups: 2,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate all architectural invariants.
|
||||
///
|
||||
/// # Errors
|
||||
@@ -108,7 +210,12 @@ impl WiFlowStdConfig {
|
||||
if self.tcn_groups == 0 {
|
||||
return Err(ConfigError::invalid_value("tcn_groups", "must be >= 1"));
|
||||
}
|
||||
if self.subcarriers % self.tcn_groups != 0 {
|
||||
// In Gcd mode the per-conv group count is gcd(channels, tcn_groups)
|
||||
// and in Depthwise mode it is the channel count itself, so the
|
||||
// divisibility invariant holds by construction; only Fixed mode
|
||||
// (upstream behavior) needs the explicit checks.
|
||||
let fixed = self.tcn_groups_mode == TcnGroupsMode::Fixed;
|
||||
if fixed && self.subcarriers % self.tcn_groups != 0 {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"subcarriers",
|
||||
format!(
|
||||
@@ -124,7 +231,7 @@ impl WiFlowStdConfig {
|
||||
));
|
||||
}
|
||||
for (i, &c) in self.tcn_channels.iter().enumerate() {
|
||||
if c == 0 || c % self.tcn_groups != 0 {
|
||||
if c == 0 || (fixed && c % self.tcn_groups != 0) {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"tcn_channels",
|
||||
format!(
|
||||
@@ -134,6 +241,18 @@ impl WiFlowStdConfig {
|
||||
));
|
||||
}
|
||||
}
|
||||
if self.input_pw_groups == 0
|
||||
|| self.subcarriers % self.input_pw_groups != 0
|
||||
|| self.tcn_channels[0] % self.input_pw_groups != 0
|
||||
{
|
||||
return Err(ConfigError::invalid_value(
|
||||
"input_pw_groups",
|
||||
format!(
|
||||
"{} must be >= 1 and divide both subcarriers={} and tcn_channels[0]={}",
|
||||
self.input_pw_groups, self.subcarriers, self.tcn_channels[0]
|
||||
),
|
||||
));
|
||||
}
|
||||
if self.conv_channels.is_empty() {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"conv_channels",
|
||||
@@ -184,19 +303,61 @@ impl WiFlowStdConfig {
|
||||
*self.tcn_channels.last().unwrap_or(&0)
|
||||
}
|
||||
|
||||
/// Width of the encoder feature map after the strided conv blocks.
|
||||
/// Group count of a grouped TCN conv over `channels` channels, per
|
||||
/// [`Self::tcn_groups_mode`].
|
||||
pub fn tcn_conv_groups(&self, channels: usize) -> usize {
|
||||
match self.tcn_groups_mode {
|
||||
TcnGroupsMode::Fixed => self.tcn_groups,
|
||||
TcnGroupsMode::Gcd => gcd(channels, self.tcn_groups),
|
||||
TcnGroupsMode::Depthwise => channels,
|
||||
}
|
||||
}
|
||||
|
||||
/// Width stride of each `AsymmetricConvBlock`, derived with the sweep's
|
||||
/// rule (`model_compact.py::compute_strides`): halve the width
|
||||
/// (`w → ceil(w / 2)`, the `(1,3)`-kernel stride-2 output size) only
|
||||
/// while the result stays ≥ [`Self::keypoints`], so the final adaptive
|
||||
/// pool never has to duplicate rows. At the upstream default
|
||||
/// (240 channels, 15 keypoints) this derives `[2, 2, 2, 2]` — the
|
||||
/// hardcoded upstream schedule, exactly.
|
||||
pub fn conv_strides(&self) -> Vec<usize> {
|
||||
let mut w = self.tcn_output_channels();
|
||||
let mut strides = Vec::with_capacity(self.conv_channels.len());
|
||||
for _ in &self.conv_channels {
|
||||
let next = w.div_ceil(2);
|
||||
if next >= self.keypoints {
|
||||
strides.push(2);
|
||||
w = next;
|
||||
} else {
|
||||
strides.push(1);
|
||||
}
|
||||
}
|
||||
strides
|
||||
}
|
||||
|
||||
/// Width of the encoder feature map after the conv blocks.
|
||||
///
|
||||
/// `ConvBlock1` preserves width; each `AsymmetricConvBlock` applies a
|
||||
/// `(1, 3)` kernel with stride `(1, 2)` and padding `(0, 1)`:
|
||||
/// `w → (w - 1) / 2 + 1`. Default: 240 → 120 → 60 → 30 → **15**.
|
||||
/// `(1, 3)` kernel with padding `(0, 1)` and the per-block stride from
|
||||
/// [`Self::conv_strides`]. Default: 240 → 120 → 60 → 30 → **15**.
|
||||
pub fn feature_width(&self) -> usize {
|
||||
let mut w = self.tcn_output_channels();
|
||||
for _ in &self.conv_channels {
|
||||
w = (w.saturating_sub(1)) / 2 + 1;
|
||||
for s in self.conv_strides() {
|
||||
if s == 2 {
|
||||
w = w.div_ceil(2);
|
||||
}
|
||||
}
|
||||
w
|
||||
}
|
||||
|
||||
/// Mid-channel count of the decoder's 3×3 conv:
|
||||
/// `max(conv_channels.last() / 2, 4)` (the sweep's floor of 4 keeps the
|
||||
/// decoder viable at very small widths; identical to the upstream `c / 2`
|
||||
/// for every channel count ≥ 8, including the default 64 → 32).
|
||||
pub fn decoder_mid(&self) -> usize {
|
||||
(self.conv_channels.last().unwrap_or(&0) / 2).max(4)
|
||||
}
|
||||
|
||||
/// Output tensor shape `(batch, keypoints, 2)`. The adaptive average pool
|
||||
/// maps the feature height to `keypoints` regardless of its size, so the
|
||||
/// keypoint count is free (15 and 17 share identical weights).
|
||||
@@ -217,10 +378,19 @@ impl WiFlowStdConfig {
|
||||
pub fn param_count(&self) -> usize {
|
||||
let mut total = 0;
|
||||
|
||||
// TCN stack.
|
||||
// TCN stack: per-conv groups follow tcn_groups_mode; only the first
|
||||
// block's pointwise/downsample convs use input_pw_groups.
|
||||
let mut c_in = self.subcarriers;
|
||||
for &c_out in &self.tcn_channels {
|
||||
total += tcn_block_params(c_in, c_out, TCN_KERNEL, self.tcn_groups);
|
||||
for (i, &c_out) in self.tcn_channels.iter().enumerate() {
|
||||
let pw_groups = if i == 0 { self.input_pw_groups } else { 1 };
|
||||
total += tcn_block_params(
|
||||
c_in,
|
||||
c_out,
|
||||
TCN_KERNEL,
|
||||
self.tcn_conv_groups(c_in),
|
||||
self.tcn_conv_groups(c_out),
|
||||
pw_groups,
|
||||
);
|
||||
c_in = c_out;
|
||||
}
|
||||
|
||||
@@ -237,8 +407,8 @@ impl WiFlowStdConfig {
|
||||
// Dual axial attention: width axis + height axis, both c_in → c_in.
|
||||
total += 2 * axial_attention_params(c_in, self.attention_groups);
|
||||
|
||||
// Decoder: 3×3 conv (c → c/2) + BN + 1×1 conv (c/2 → 2) + BN.
|
||||
total += decoder_params(c_in);
|
||||
// Decoder: 3×3 conv (c → decoder_mid) + BN + 1×1 conv (mid → 2) + BN.
|
||||
total += decoder_params(c_in, self.decoder_mid());
|
||||
|
||||
total
|
||||
}
|
||||
@@ -250,18 +420,29 @@ impl WiFlowStdConfig {
|
||||
|
||||
/// One `InnerGroupedTemporalBlock`: two (depthwise-grouped conv → BN →
|
||||
/// pointwise conv → BN) stages plus a 1×1 + BN residual projection when the
|
||||
/// channel count changes. All convs are bias-free.
|
||||
fn tcn_block_params(c_in: usize, c_out: usize, k: usize, groups: usize) -> usize {
|
||||
let grouped1 = c_in * (c_in / groups) * k; // depthwise-grouped, c_in → c_in
|
||||
/// channel count changes. All convs are bias-free. `g_in`/`g_out` are the
|
||||
/// group counts of the two grouped convs (each conv groups over its own
|
||||
/// channel count — they differ in `Gcd`/`Depthwise` mode); `pw_groups`
|
||||
/// groups the first pointwise conv and the residual projection (the sweep's
|
||||
/// `input_pw_groups`, block 0 only — 1 everywhere else).
|
||||
fn tcn_block_params(
|
||||
c_in: usize,
|
||||
c_out: usize,
|
||||
k: usize,
|
||||
g_in: usize,
|
||||
g_out: usize,
|
||||
pw_groups: usize,
|
||||
) -> usize {
|
||||
let grouped1 = c_in * (c_in / g_in) * k; // depthwise-grouped, c_in → c_in
|
||||
let bn1g = 2 * c_in;
|
||||
let pw1 = c_out * c_in; // pointwise 1×1
|
||||
let pw1 = c_out * (c_in / pw_groups); // pointwise 1×1
|
||||
let bn1p = 2 * c_out;
|
||||
let grouped2 = c_out * (c_out / groups) * k;
|
||||
let grouped2 = c_out * (c_out / g_out) * k;
|
||||
let bn2g = 2 * c_out;
|
||||
let pw2 = c_out * c_out;
|
||||
let bn2p = 2 * c_out;
|
||||
let downsample = if c_in != c_out {
|
||||
c_in * c_out + 2 * c_out
|
||||
(c_in / pw_groups) * c_out + 2 * c_out
|
||||
} else {
|
||||
0
|
||||
};
|
||||
@@ -288,10 +469,9 @@ fn axial_attention_params(c: usize, groups: usize) -> usize {
|
||||
qkv + bn_qkv + bn_similarity + bn_output
|
||||
}
|
||||
|
||||
/// Decoder: `Conv2d(c → c/2, 3×3, bias)` + BN + `Conv2d(c/2 → 2, 1×1, bias)`
|
||||
/// + BN.
|
||||
fn decoder_params(c: usize) -> usize {
|
||||
let mid = c / 2;
|
||||
/// Decoder: `Conv2d(c → mid, 3×3, bias)` + BN + `Conv2d(mid → 2, 1×1, bias)`
|
||||
/// + BN, where `mid` = [`WiFlowStdConfig::decoder_mid`].
|
||||
fn decoder_params(c: usize, mid: usize) -> usize {
|
||||
let conv1 = mid * c * 9 + mid;
|
||||
let bn1 = 2 * mid;
|
||||
let conv2 = 2 * mid + 2;
|
||||
@@ -335,10 +515,10 @@ mod tests {
|
||||
#[test]
|
||||
fn per_component_breakdown_matches_hand_calculation() {
|
||||
// TCN levels (hand-verified against the reference layer shapes).
|
||||
assert_eq!(tcn_block_params(540, 540, 3, 20), 675_000);
|
||||
assert_eq!(tcn_block_params(540, 440, 3, 20), 746_180);
|
||||
assert_eq!(tcn_block_params(440, 340, 3, 20), 464_780);
|
||||
assert_eq!(tcn_block_params(340, 240, 3, 20), 249_380);
|
||||
assert_eq!(tcn_block_params(540, 540, 3, 20, 20, 1), 675_000);
|
||||
assert_eq!(tcn_block_params(540, 440, 3, 20, 20, 1), 746_180);
|
||||
assert_eq!(tcn_block_params(440, 340, 3, 20, 20, 1), 464_780);
|
||||
assert_eq!(tcn_block_params(340, 240, 3, 20, 20, 1), 249_380);
|
||||
// Conv encoder.
|
||||
assert_eq!(conv_block_params(1, 8), 504);
|
||||
assert_eq!(conv_block_params(8, 8), 728);
|
||||
@@ -347,7 +527,131 @@ mod tests {
|
||||
assert_eq!(conv_block_params(32, 64), 33_472);
|
||||
// Attention + decoder.
|
||||
assert_eq!(axial_attention_params(64, 8), 12_816);
|
||||
assert_eq!(decoder_params(64), 18_598);
|
||||
assert_eq!(decoder_params(64, 32), 18_598);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// ADR-152 efficiency-sweep compact presets. The parameter pins are
|
||||
// GROUND TRUTH measured from the trained PyTorch checkpoints
|
||||
// (benchmarks/wiflow-std/results/efficiency_sweep.jsonl, 2026-06-11):
|
||||
// any mismatch means the Rust formula or config mapping is wrong.
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn half_preset_param_count_matches_trained_checkpoint() {
|
||||
let cfg = WiFlowStdConfig::half();
|
||||
cfg.validate().expect("half preset must validate");
|
||||
assert_eq!(cfg.param_count(), 843_834);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quarter_preset_param_count_matches_trained_checkpoint() {
|
||||
let cfg = WiFlowStdConfig::quarter();
|
||||
cfg.validate().expect("quarter preset must validate");
|
||||
assert_eq!(cfg.param_count(), 338_600);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_preset_param_count_matches_trained_checkpoint() {
|
||||
let cfg = WiFlowStdConfig::tiny();
|
||||
cfg.validate().expect("tiny preset must validate");
|
||||
assert_eq!(cfg.param_count(), 56_290);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preset_tcn_groups_match_sweep_per_block_record() {
|
||||
// efficiency_sweep.jsonl "tcn_groups_per_block": (conv1, conv2) of
|
||||
// each block — conv1 groups over c_in, conv2 over c_out.
|
||||
let half = WiFlowStdConfig::half();
|
||||
let groups: Vec<(usize, usize)> = {
|
||||
let mut c_in = half.subcarriers;
|
||||
half.tcn_channels
|
||||
.iter()
|
||||
.map(|&c_out| {
|
||||
let g = (half.tcn_conv_groups(c_in), half.tcn_conv_groups(c_out));
|
||||
c_in = c_out;
|
||||
g
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
assert_eq!(groups, [(20, 10), (10, 20), (20, 10), (10, 20)]);
|
||||
|
||||
let tiny = WiFlowStdConfig::tiny();
|
||||
assert_eq!(tiny.tcn_conv_groups(540), 540); // depthwise input conv
|
||||
assert_eq!(tiny.tcn_conv_groups(68), 68);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preset_stride_schedules_match_sweep_record() {
|
||||
// efficiency_sweep.jsonl "conv_strides" / "final_width".
|
||||
assert_eq!(WiFlowStdConfig::default().conv_strides(), [2, 2, 2, 2]);
|
||||
assert_eq!(WiFlowStdConfig::half().conv_strides(), [2, 2, 2, 1]);
|
||||
assert_eq!(WiFlowStdConfig::quarter().conv_strides(), [2, 2, 1, 1]);
|
||||
assert_eq!(WiFlowStdConfig::tiny().conv_strides(), [2, 1, 1, 1]);
|
||||
assert_eq!(WiFlowStdConfig::half().feature_width(), 15);
|
||||
assert_eq!(WiFlowStdConfig::quarter().feature_width(), 15);
|
||||
assert_eq!(WiFlowStdConfig::tiny().feature_width(), 16);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fixed_mode_with_defaults_is_unchanged_by_new_knobs() {
|
||||
// The new fields default to upstream behavior: gcd(c, 20) == 20 for
|
||||
// every default channel count, so Gcd mode is also a no-op there.
|
||||
let mut cfg = WiFlowStdConfig::default();
|
||||
assert_eq!(cfg.param_count(), REFERENCE_PARAMS);
|
||||
cfg.tcn_groups_mode = TcnGroupsMode::Gcd;
|
||||
cfg.validate().expect("gcd mode validates at defaults");
|
||||
assert_eq!(cfg.param_count(), REFERENCE_PARAMS);
|
||||
assert_eq!(WiFlowStdConfig::default().decoder_mid(), 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_bad_input_pw_groups() {
|
||||
// 7 divides neither 540 nor 540's first TCN level.
|
||||
let cfg = WiFlowStdConfig {
|
||||
input_pw_groups: 7,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(cfg.validate().is_err());
|
||||
// 27 divides subcarriers=540 but not tiny's tcn_channels[0]=68.
|
||||
let cfg = WiFlowStdConfig {
|
||||
input_pw_groups: 27,
|
||||
..WiFlowStdConfig::tiny()
|
||||
};
|
||||
assert!(cfg.validate().is_err());
|
||||
let zero = WiFlowStdConfig {
|
||||
input_pw_groups: 0,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(zero.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serde_defaults_for_new_fields_are_backward_compatible() {
|
||||
// A config serialized before the compact-variant knobs existed must
|
||||
// deserialize to upstream behavior (Fixed mode, input_pw_groups 1).
|
||||
let legacy = r#"{
|
||||
"subcarriers": 540, "window": 20,
|
||||
"tcn_channels": [540, 440, 340, 240], "tcn_groups": 20,
|
||||
"conv_channels": [8, 16, 32, 64], "attention_groups": 8,
|
||||
"keypoints": 15, "dropout": 0.5
|
||||
}"#;
|
||||
let cfg: WiFlowStdConfig = serde_json::from_str(legacy).expect("deserialize");
|
||||
assert_eq!(cfg, WiFlowStdConfig::default());
|
||||
assert_eq!(cfg.param_count(), REFERENCE_PARAMS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serde_roundtrip_preserves_presets() {
|
||||
for cfg in [
|
||||
WiFlowStdConfig::half(),
|
||||
WiFlowStdConfig::quarter(),
|
||||
WiFlowStdConfig::tiny(),
|
||||
] {
|
||||
let json = serde_json::to_string(&cfg).expect("serialize");
|
||||
let back: WiFlowStdConfig = serde_json::from_str(&json).expect("deserialize");
|
||||
assert_eq!(back, cfg);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -41,12 +41,20 @@ pub(super) struct GroupedTemporalBlock {
|
||||
}
|
||||
|
||||
impl GroupedTemporalBlock {
|
||||
/// `g_in`/`g_out`: group counts of the two grouped convs (each conv
|
||||
/// groups over its own channel count — they differ under the ADR-152
|
||||
/// compact variants' `Gcd`/`Depthwise` modes). `pw_groups` groups the
|
||||
/// first pointwise conv and the residual projection (`input_pw_groups`
|
||||
/// on block 0; 1 everywhere else).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn new(
|
||||
vs: nn::Path,
|
||||
c_in: i64,
|
||||
c_out: i64,
|
||||
dilation: i64,
|
||||
groups: i64,
|
||||
g_in: i64,
|
||||
g_out: i64,
|
||||
pw_groups: i64,
|
||||
dropout: f64,
|
||||
) -> Self {
|
||||
let k = TCN_KERNEL as i64;
|
||||
@@ -58,24 +66,25 @@ impl GroupedTemporalBlock {
|
||||
bias: false,
|
||||
..Default::default()
|
||||
};
|
||||
let pointwise_cfg = nn::ConvConfig {
|
||||
let pointwise_cfg = |groups| nn::ConvConfig {
|
||||
groups,
|
||||
bias: false,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let conv1_group = nn::conv1d(&vs / "conv1_group", c_in, c_in, k, grouped_cfg(groups));
|
||||
let conv1_group = nn::conv1d(&vs / "conv1_group", c_in, c_in, k, grouped_cfg(g_in));
|
||||
let bn1_group = nn::batch_norm1d(&vs / "bn1_group", c_in, bn_cfg());
|
||||
let conv1_pw = nn::conv1d(&vs / "conv1_pw", c_in, c_out, 1, pointwise_cfg);
|
||||
let conv1_pw = nn::conv1d(&vs / "conv1_pw", c_in, c_out, 1, pointwise_cfg(pw_groups));
|
||||
let bn1_pw = nn::batch_norm1d(&vs / "bn1_pw", c_out, bn_cfg());
|
||||
|
||||
let conv2_group = nn::conv1d(&vs / "conv2_group", c_out, c_out, k, grouped_cfg(groups));
|
||||
let conv2_group = nn::conv1d(&vs / "conv2_group", c_out, c_out, k, grouped_cfg(g_out));
|
||||
let bn2_group = nn::batch_norm1d(&vs / "bn2_group", c_out, bn_cfg());
|
||||
let conv2_pw = nn::conv1d(&vs / "conv2_pw", c_out, c_out, 1, pointwise_cfg);
|
||||
let conv2_pw = nn::conv1d(&vs / "conv2_pw", c_out, c_out, 1, pointwise_cfg(1));
|
||||
let bn2_pw = nn::batch_norm1d(&vs / "bn2_pw", c_out, bn_cfg());
|
||||
|
||||
let downsample = (c_in != c_out).then(|| {
|
||||
(
|
||||
nn::conv1d(&vs / "ds_conv", c_in, c_out, 1, pointwise_cfg),
|
||||
nn::conv1d(&vs / "ds_conv", c_in, c_out, 1, pointwise_cfg(pw_groups)),
|
||||
nn::batch_norm1d(&vs / "ds_bn", c_out, bn_cfg()),
|
||||
)
|
||||
});
|
||||
|
||||
@@ -63,7 +63,7 @@ mod layers;
|
||||
#[cfg(feature = "tch-backend")]
|
||||
pub mod model;
|
||||
|
||||
pub use config::WiFlowStdConfig;
|
||||
pub use config::{TcnGroupsMode, WiFlowStdConfig};
|
||||
|
||||
#[cfg(feature = "tch-backend")]
|
||||
pub use model::WiFlowStdModel;
|
||||
|
||||
@@ -59,33 +59,40 @@ impl WiFlowStdModel {
|
||||
let vs = nn::VarStore::new(device);
|
||||
let root = vs.root();
|
||||
|
||||
// TCN stack: dilation doubles per level, causal padding.
|
||||
// TCN stack: dilation doubles per level, causal padding. Per-conv
|
||||
// groups follow `config.tcn_groups_mode`; only block 0's pointwise/
|
||||
// downsample convs use `config.input_pw_groups` (ADR-152 sweep).
|
||||
let mut tcn = Vec::with_capacity(config.tcn_channels.len());
|
||||
let mut c_in = config.subcarriers as i64;
|
||||
let mut c_in = config.subcarriers;
|
||||
for (i, &c_out) in config.tcn_channels.iter().enumerate() {
|
||||
let dilation = 1_i64 << i;
|
||||
let pw_groups = if i == 0 { config.input_pw_groups } else { 1 };
|
||||
tcn.push(GroupedTemporalBlock::new(
|
||||
&root / format!("tcn{i}"),
|
||||
c_in,
|
||||
c_in as i64,
|
||||
c_out as i64,
|
||||
dilation,
|
||||
config.tcn_groups as i64,
|
||||
config.tcn_conv_groups(c_in) as i64,
|
||||
config.tcn_conv_groups(c_out) as i64,
|
||||
pw_groups as i64,
|
||||
config.dropout,
|
||||
));
|
||||
c_in = c_out as i64;
|
||||
c_in = c_out;
|
||||
}
|
||||
|
||||
// 2-D conv encoder: ConvBlock1 (stride 1) + strided asymmetric blocks.
|
||||
// 2-D conv encoder: ConvBlock1 (stride 1) + asymmetric blocks with
|
||||
// the derived stride schedule ([2, 2, 2, 2] at the upstream default).
|
||||
let c0 = config.conv_channels[0] as i64;
|
||||
let conv_in = ConvBlock::new(&root / "conv_in", 1, c0, 1);
|
||||
let mut conv_blocks = Vec::with_capacity(config.conv_channels.len());
|
||||
let strides = config.conv_strides();
|
||||
let mut c_in = c0;
|
||||
for (i, &c_out) in config.conv_channels.iter().enumerate() {
|
||||
conv_blocks.push(ConvBlock::new(
|
||||
&root / format!("conv{i}"),
|
||||
c_in,
|
||||
c_out as i64,
|
||||
2,
|
||||
strides[i] as i64,
|
||||
));
|
||||
c_in = c_out as i64;
|
||||
}
|
||||
@@ -93,8 +100,8 @@ impl WiFlowStdModel {
|
||||
let attention =
|
||||
DualAxialAttention::new(&root / "attention", c_in, config.attention_groups as i64);
|
||||
|
||||
// Decoder: c → c/2 (3×3) → 2 (1×1), BN + SiLU after each conv.
|
||||
let mid = c_in / 2;
|
||||
// Decoder: c → decoder_mid (3×3) → 2 (1×1), BN + SiLU after each conv.
|
||||
let mid = config.decoder_mid() as i64;
|
||||
let dec_conv1 = nn::conv2d(
|
||||
&root / "dec_conv1",
|
||||
c_in,
|
||||
@@ -241,6 +248,26 @@ mod tests {
|
||||
assert_eq!(model.num_parameters(), 2_225_042);
|
||||
}
|
||||
|
||||
/// ADR-152 efficiency-sweep compact presets: the tch graph must realise
|
||||
/// exactly the trained checkpoints' measured parameter counts
|
||||
/// (benchmarks/wiflow-std/results/efficiency_sweep.jsonl) and produce
|
||||
/// the standard [B, 15, 2] output.
|
||||
#[test]
|
||||
fn compact_preset_param_counts_and_shapes() {
|
||||
for (cfg, expected) in [
|
||||
(WiFlowStdConfig::half(), 843_834_i64),
|
||||
(WiFlowStdConfig::quarter(), 338_600),
|
||||
(WiFlowStdConfig::tiny(), 56_290),
|
||||
] {
|
||||
tch::manual_seed(0);
|
||||
let model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("preset builds");
|
||||
assert_eq!(model.num_parameters(), expected);
|
||||
assert_eq!(model.num_parameters(), cfg.param_count() as i64);
|
||||
let out = model.forward_inference(&random_csi(&cfg, 2));
|
||||
assert_eq!(out.size(), &[2, 15, 2]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_output_shape_15_keypoints() {
|
||||
tch::manual_seed(0);
|
||||
|
||||
Reference in New Issue
Block a user