diff --git a/benchmarks/wiflow-std/RESULTS.md b/benchmarks/wiflow-std/RESULTS.md index a761216c..85820340 100644 --- a/benchmarks/wiflow-std/RESULTS.md +++ b/benchmarks/wiflow-std/RESULTS.md @@ -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 diff --git a/benchmarks/wiflow-std/results/edge_optimization.json b/benchmarks/wiflow-std/results/edge_optimization.json index ba07eac5..dfc0e143 100644 --- a/benchmarks/wiflow-std/results/edge_optimization.json +++ b/benchmarks/wiflow-std/results/edge_optimization.json @@ -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 + } + } } } \ No newline at end of file diff --git a/benchmarks/wiflow-std/tiny_edge_bench.py b/benchmarks/wiflow-std/tiny_edge_bench.py new file mode 100644 index 00000000..3d995d9b --- /dev/null +++ b/benchmarks/wiflow-std/tiny_edge_bench.py @@ -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 ] [--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() diff --git a/v2/crates/wifi-densepose-train/src/wiflow_std/config.rs b/v2/crates/wifi-densepose-train/src/wiflow_std/config.rs index ed4e020b..2068a47d 100644 --- a/v2/crates/wifi-densepose-train/src/wiflow_std/config.rs +++ b/v2/crates/wifi-densepose-train/src/wiflow_std/config.rs @@ -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 { + 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] diff --git a/v2/crates/wifi-densepose-train/src/wiflow_std/layers.rs b/v2/crates/wifi-densepose-train/src/wiflow_std/layers.rs index 5c514030..2bec866f 100644 --- a/v2/crates/wifi-densepose-train/src/wiflow_std/layers.rs +++ b/v2/crates/wifi-densepose-train/src/wiflow_std/layers.rs @@ -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()), ) }); diff --git a/v2/crates/wifi-densepose-train/src/wiflow_std/mod.rs b/v2/crates/wifi-densepose-train/src/wiflow_std/mod.rs index e927d73a..0ff360f4 100644 --- a/v2/crates/wifi-densepose-train/src/wiflow_std/mod.rs +++ b/v2/crates/wifi-densepose-train/src/wiflow_std/mod.rs @@ -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; diff --git a/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs b/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs index 9e2665b3..6d0bf813 100644 --- a/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs +++ b/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs @@ -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);