mirror of
https://github.com/ruvnet/RuView
synced 2026-08-08 20:11:43 +00:00
ADR-152: WiFi-Pose SOTA 2026 intake — WiFlow-STD benchmark, Rust integrations, ADR-153 802.11bf layer, efficiency frontier (#1008)
* feat(calibration): NodeGeometry transceiver-geometry recording (ADR-152 §2.1.1) PerceptAlign-motivated geometry capture at enrollment: per-node optional records (position, antenna orientation, inter-node distances, acquisition method) — recorded when known, never required. Event-sourced via EnrollmentEvent::GeometryRecorded (latest recording wins); persisted on SpecialistBank with serde defaults so pre-ADR-152 bank JSON loads cleanly (fixture-proven, and geometry-free banks serialize byte-shape-identical to the old schema); threaded through MultiNodeMixture as data only — the learned geometry embeddings and algorithmic fusion use are §2.1.2, deliberately deferred until the ADR-151 P6 LoRA heads exist. Geometry recorded from now on means banks captured today remain usable for layout-conditioned training later — you can't retroactively add geometry to data you didn't record. 8 new tests (3 geometry, 2 anchor, 2 bank, 1 multistatic) + full-loop extension (2-node geometry, one tape-measured + one unknown, surviving the bank JSON round-trip the runtime loads from). 50/50 calibration (both feature configs) + 23 CLI tests green. Co-Authored-By: RuFlo <ruv@ruv.net> * feat(training): two-checkerboard camera↔room calibration for ADR-079 labels (ADR-152 §2.1.3) Defends the camera-supervised pipeline against PerceptAlign's "coordinate overfitting": MediaPipe keypoints were emitted in raw camera coordinates with no shared frame and no transceiver-geometry metadata — the exact label shape that memorizes deployment layout and collapses cross-layout. - scripts/calibrate-camera-room.py + calibration_lib.py: OpenCV two-checkerboard calibration → versioned bundle JSON (intrinsics, camera→room extrinsics, checkerboard spec, transceiver geometry, sha256 calibration_id). Intrinsics resolve from file > cache > multi-view computation > loud-warning 2-view fallback. - collect-ground-truth.py --calibration <bundle>: every sample gains keypoints_room (unit bearing rays from the camera center in the room frame — documented projective alignment; raw image coords preserved so training chooses), camera_origin_room, calibration_id, and the transceiver geometry stamp. Without the flag, output is byte-identical to before (tested) + a one-line ADR-152 warning. Design finding (recorded for ADR-152): a single planar checkerboard's corner grid is centrosymmetric — the reversed corner ordering fits a ghost camera pose with IDENTICAL reprojection error, so per-board flip disambiguation is mathematically ill-posed. solve_two_board_extrinsics solves the joint wall+floor set over all 4 flip combinations, where the minimum is unique — an independent reason the TWO-checkerboard method is required, beyond what PerceptAlign states. 15 headless pytest tests green (synthetic corners: extrinsics recovery incl. ghost resolution, bundle round-trip + hash stability, ray transforms w/ distortion + cross-resolution, no-calibration byte identity). Co-Authored-By: RuFlo <ruv@ruv.net> * feat(benchmarks): WiFlow-STD reproduction harness + measurement (a) results (ADR-152 §2.2) Shipped checkpoint REFUTED (0.08% PCK@20, wrong keypoint normalization); 6 reproducibility defects documented (broken imports, corrupted dataset tail with float32-max garbage that NaN-poisons fp16 BatchNorm, unreachable test phase). After repairs, retraining with upstream defaults reproduces 96.09% PCK@20 full-test / 96.61% corruption-free (published 97.25%) on RTX 5080. Claims graded MEASURED-EQUIVALENT; 2.23M params + ~0.055 GFLOPs verified. Third-party code/weights/data stay out of tree (gitignored). Co-Authored-By: claude-flow <ruv@ruv.net> * feat: ADR-152 Rust integrations + ADR-153 802.11bf protocol model - calibration: GeometryEmbedding — 32-slot permutation-invariant NodeGeometry featurization for future LoRA-head conditioning (ADR-152 §2.1.2); derived SpecialistBank::geometry_embedding() accessor; 59 tests - train: MaePretrainConfig + patchify/random-mask with UNSW measured recipe (80% masking, (30,3) patches; ADR-152 §2.3, arXiv 2511.18792); strict no-truncate/no-NaN policy; proptest properties - train: WiFlowStdModel — tch-gated port of the verified ~96%-PCK@20 WiFlow-STD architecture (ADR-152 §2.2 beyond-SOTA); ungated param formula pinned to 2,225,042; 15/17-keypoint support; 239 crate tests - hardware: ieee80211bf forward-compatibility protocol model (ADR-153): SpecProfile gates, SensingCapabilities negotiation, required ConsentMode, session FSM, SensingTransport + SimTransport + OpportunisticCsiBridge; full acceptance checklist covered; 156+4 tests - deps: ruvector bumps per ADR-152 §2.6 survey (mincut/solver 2.0.6, attention 2.1.0, gnn 2.2.0); vendor/ruvector synced to a083bd77f - docs: ADR-153 accepted; ADR-152 §2.2 status, §2.4 amendment, §2.6 added Workspace: 162 test suites green (--no-default-features); Python proof PASS. Known pre-existing flake: homecore-api env_empty_falls_back_to_defaults (unserialized env-var mutation) — untouched, follow-up. Co-Authored-By: claude-flow <ruv@ruv.net> * docs: CHANGELOG + CLAUDE.md entries for ADR-152 integrations and ADR-153 Co-Authored-By: claude-flow <ruv@ruv.net> * fix(train): repair tch-backend bit-rot — gated path compiles and tests run again Mechanical API refresh against current tch: Vec::from(Tensor) -> try_from (+ explicit flatten), numel() usize cast, Rem/div ops -> remainder() / divide_scalar_mode(floor) — the latter fixed a silent true-division bug in heatmap argmax decoding; clamp(1.0, f64::MAX) -> clamp_min (torch 2.x scalar overflow panic); petgraph EdgeRef import; missing EvalMetrics and verify_checkpoint_dir APIs that tests documented. wiflow_std roundtrip test uses safetensors (.pt _save_parameters roundtrip broken in torch 2.11 Windows). Gated: 349 passed (incl. all 20 wiflow_std); ungated: unchanged. Known pre-existing: gaussian-heatmap convention mismatch (2 tests), proof seed race under parallel threads — documented, deliberate follow-ups. Co-Authored-By: claude-flow <ruv@ruv.net> * feat(train): WiFlow-STD PyTorch->tch weight import + numerical parity proof export_to_safetensors.py maps the retrained checkpoint (295 tensors -> 248 mapped, param sum exactly 2,225,042; num_batches_tracked dropped) into a tch-loadable safetensors plus a deterministic parity fixture. Gated #[ignore] integration test loads it strictly and asserts forward-pass agreement: max abs diff 1.192e-7 on the seed-42 fixture. dump_variable_names test makes the tch name layout authoritative. Zero architecture discrepancies found. Co-Authored-By: claude-flow <ruv@ruv.net> * fix: workflow-review findings — BN gamma init, ThresholdParams serde, init docs Concurrent validation workflow (2 review lanes + adversarial verification, 13 agents): 5 confirmed findings, 3 refuted. Fixes: - wiflow_std: pin BatchNorm gamma to 1.0 (tch default draws Uniform(0,1) — silently halves activations in from-scratch training; loaded checkpoints unaffected, parity re-verified after the change) - wiflow_std: document the conv-init divergences vs the reference's effective kaiming_normal(fan_out) re-init (from-scratch dynamics only) - ieee80211bf: ThresholdParams deserialization validates via try_from so the <=100 invariant holds for untrusted payloads (+ rejection test) Benchmarks (release, ruvzen): GeometryEmbedding 1.84us/call (542k/s), MAE tokenization 7.38us/window (135k/s), 802.11bf FSM 8.9M events/s — nothing suspicious. Co-Authored-By: claude-flow <ruv@ruv.net> * docs(adr): ADR-152 §2.1.4 gate resolved — PerceptAlign repo MIT, dataset on HF Co-Authored-By: claude-flow <ruv@ruv.net> * feat(benchmarks): edge optimization measured + measurement (b) blocked + 92.9% retraction Edge optimization (ADR-152 optimize track): ONNX Runtime fp32 is the CPU latency win (3.2 ms/window, ~3.4x faster than torch, parity 2.4e-7); ORT dynamic int8 reaches 2.44 MB (paper's ~2.2 MB claim plausible only via conv-capable toolchains; -0.16pt PCK@20, +18% MPJPE, 2x slower); torch dynamic quant converts 0% of this conv-only model; fp16 halves storage free but is slower on CPU. Measurement (b) BLOCKED-ON-DATA: only 1,077 paired ESP32 windows exist (stop rule <2k). Forensic recheck of the surviving April holdout RETRACTS the ADR-079 '92.9% PCK@20' figure: constant-output model, absolute (not torso) threshold, 69 near-static frames — mean predictor scores 100% under that protocol; torso-PCK@20 is 19.1%. Corroborates PR #535. Stale citations removed from user-guide, readme-details, ADR-152 §2.1.3; no-citation rule extended to ADR-079 accuracy claims. Unblock: >=2k-window multi-pose paired session + torso-PCK re-baseline. Co-Authored-By: claude-flow <ruv@ruv.net> * docs(user-guide): corrected camera-supervised collection tutorial Step 0 CSI-rate check + session-length math (window yield = frames/20 — the May session's 8x under-delivery was a ~12 Hz CSI rate, not an aligner bug); two-checkerboard calibration step (ADR-152 §2.1.3); pose-variety and confidence guidance; torso-normalized PCK + temporal-split + pred-variance eval protocol (lessons from the 92.9% retraction); scale presets re-keyed to realistic window counts. Co-Authored-By: claude-flow <ruv@ruv.net> * feat(benchmarks): static PTQ int8 (calibrated) results + overnight capture script Conv-only static QDQ beats dynamic int8 on accuracy (PCK@20 96.61-96.63% vs 96.52%, MPJPE +10% vs +18% over fp32) at ~equal size/latency; all-ops QDQ strictly worse (int8 activations through attention glue). Entropy calibration verified bit-identical to MinMax on this data. Deployment: ONNX fp32 for speed (3.2ms), static conv-only QDQ for smallest (2.53MB). Also: scripts/overnight-empty-capture.py — segmented UDP CSI recorder for empty-room baselines (no glob collisions, detach-safe). Co-Authored-By: claude-flow <ruv@ruv.net> * feat(benchmarks): measurement (b) MEASURED — optimization transfer only, mean-pose baseline wins WiFlow-STD fine-tuned on 2,046 fresh single-room ESP32 paired windows (temporal 70/15/15, 70->540 adapter, K=17): pretrained-init 65% PCK@20 vs scratch 0% (optimization transfer) but frozen-trunk ~0% (no feature transfer), and NOTHING beats the mean-pose baseline (95.9% PCK@20 — single subject, near-static normalized coords). Honesty gates held: pred std 0.0113 (non-constant model) but mean-baseline dominance means no citable CSI->pose capability from this data. ADR-152 open question 1 answered partially; definitive answer needs multi-subject/position data. Two new aligner findings: heterogeneous csi_shape with silent zero-padding (~20%), and extractCsiMatrix's transposed shape label (frame-major data, [nSc, nFrames] label) — fixes pending. Co-Authored-By: claude-flow <ruv@ruv.net> * feat(benchmarks): efficiency sweep MEASURED — half model dominates full reference Compact WiFlow-STD variants on the same data/split/protocol: half (843,834 params, 0.38x) strictly dominates the 2.23M reference (PCK@20 96.62 vs 96.61, PCK@50 99.47 vs 99.11, MPJPE 0.00898 vs 0.0094) — the published architecture is over-parameterized for its own benchmark. quarter (338k) 96.05%; tiny (56,290 params, 1/39.5) holds 94.11% — a ~220KB fp32 edge candidate. In-domain caveats recorded; cross-domain untested. Co-Authored-By: claude-flow <ruv@ruv.net> * 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> * fix: resolve all 10 confirmed code-review findings (7-angle review, 20/20 verified) wiflow_std: min_feature_width (default 15) replaces the keypoints->stride coupling — for_keypoints(17) now provably builds the trained [2,2,2,2] graph and pools 15->17, matching the validated Python protocol (pinned by tests); param_count() total on invalid configs; random_mask returns Result and rejects non-finite/out-of-range ratios; trainer checkpoints switched to safetensors (.pt VarStore roundtrip broken on Windows torch 2.11). ieee80211bf: SBP proxy now re-triggers instances and relays reports via Action::RelaySbpReport -> SensingFrame::SbpReport (clients consume via their existing path); missed_instances reset on success = consecutive semantics; SessionTable gains a guarded SBP entry point + unknown-id drop counter; initiator-role sessions reject inbound setup/SBP requests (RejectedNotSupported) closing the idle hijack; StartSetup/StartSbp outside Idle return InvalidStateForCommand; SBP validation unified through evaluate_setup with a 1:1 SetupStatus->SbpStatus mapping. events.rs split out to honor the 500-line cap. calibration/cli: enrollment geometry now actually reaches trained banks — both production call sites attach .with_geometry; --geometry flag on train-room and POST /enroll/geometry + train-body geometry on calibrate-serve give production a recording surface; geometry-free banks log the ADR-152 §2.1.2 note. benchmarks: corruption masks committed as ground truth (unregenerable after in-place cleaning; verified bit-identical regeneration from the pristine copy) + generate_corruption_masks.py producer; _bench_common.py dedups the 5x-copied shim/evaluate/seed/remap (post-refactor PCK@20 re-verified equal to the last digit); remote scripts get the mmap patch; tiny_edge --calib validated multiple-of-64; onnx_bench --help no longer executes (and overwrote) the export — artifact restored byte-exact. Workspace: 2,963 tests passed, 0 failed; Python proof PASS. Co-Authored-By: claude-flow <ruv@ruv.net> * ci: build workspace tests without debuginfo — runner disk exhaustion The combined 38-crate debug target exceeds the GitHub runner's disk ('final link failed: No space left on device'); the same tree measured 151GB locally with full debuginfo. CARGO_PROFILE_{DEV,TEST}_DEBUG=0 shrinks the target ~5-10x; debuginfo serves no purpose in CI test runs. Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
rUv
@@ -11,7 +11,8 @@
|
||||
//! TrainError (top-level)
|
||||
//! ├── ConfigError (config validation / file loading)
|
||||
//! ├── DatasetError (data loading, I/O, format)
|
||||
//! └── SubcarrierError (frequency-axis resampling)
|
||||
//! ├── SubcarrierError (frequency-axis resampling)
|
||||
//! └── MaeError (MAE patchify / masking — ADR-152 §2.3)
|
||||
//! ```
|
||||
|
||||
use std::path::PathBuf;
|
||||
@@ -44,6 +45,10 @@ pub enum TrainError {
|
||||
#[error("Dataset error: {0}")]
|
||||
Dataset(#[from] DatasetError),
|
||||
|
||||
/// A MAE pretraining patchify / masking error (ADR-152 §2.3).
|
||||
#[error("MAE pretraining error: {0}")]
|
||||
Mae(#[from] MaeError),
|
||||
|
||||
/// JSON (de)serialization error.
|
||||
#[error("JSON error: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
@@ -373,3 +378,85 @@ impl SubcarrierError {
|
||||
SubcarrierError::NumericalError(msg.into())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// MaeError
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Errors produced by the MAE pretraining patchify / masking functions
|
||||
/// ([`crate::mae`], ADR-152 §2.3).
|
||||
#[derive(Debug, Error)]
|
||||
pub enum MaeError {
|
||||
/// The flat window buffer does not match the declared `time × subc` shape.
|
||||
#[error(
|
||||
"Window length {actual} does not match time × subcarriers = \
|
||||
{time} × {subc} = {expected}"
|
||||
)]
|
||||
WindowShapeMismatch {
|
||||
/// Declared time dimension.
|
||||
time: usize,
|
||||
/// Declared subcarrier dimension.
|
||||
subc: usize,
|
||||
/// Expected buffer length (`time * subc`).
|
||||
expected: usize,
|
||||
/// Actual buffer length.
|
||||
actual: usize,
|
||||
},
|
||||
|
||||
/// A patch dimension is larger than the window along that axis.
|
||||
#[error("Patch {axis} extent {patch} exceeds window {axis} extent {window}")]
|
||||
PatchExceedsWindow {
|
||||
/// Axis name (`"time"` or `"subcarrier"`).
|
||||
axis: &'static str,
|
||||
/// Patch extent along the axis.
|
||||
patch: usize,
|
||||
/// Window extent along the axis.
|
||||
window: usize,
|
||||
},
|
||||
|
||||
/// The window is not an exact multiple of the patch extent along an axis.
|
||||
///
|
||||
/// Patchification never silently truncates; crop the window to `crop`
|
||||
/// (the largest divisible extent) or change the patch size.
|
||||
#[error(
|
||||
"Window {axis} extent {window} is not divisible by patch {axis} extent \
|
||||
{patch} (remainder {remainder}); crop the window to {crop} or change \
|
||||
the patch size"
|
||||
)]
|
||||
NotDivisible {
|
||||
/// Axis name (`"time"` or `"subcarrier"`).
|
||||
axis: &'static str,
|
||||
/// Window extent along the axis.
|
||||
window: usize,
|
||||
/// Patch extent along the axis.
|
||||
patch: usize,
|
||||
/// `window % patch`.
|
||||
remainder: usize,
|
||||
/// Largest divisible extent (`window - remainder`).
|
||||
crop: usize,
|
||||
},
|
||||
|
||||
/// The mask ratio is not a finite value strictly inside `(0, 1)` — the
|
||||
/// same rule as [`MaePretrainConfig::validate`]. A NaN ratio must never
|
||||
/// silently mask zero patches, and ratios ≤ 0 / ≥ 1 degenerate to
|
||||
/// all-visible / all-masked grids.
|
||||
///
|
||||
/// [`MaePretrainConfig::validate`]: crate::mae::MaePretrainConfig::validate
|
||||
#[error("Invalid mask ratio {ratio}: must be finite and strictly inside (0, 1)")]
|
||||
InvalidMaskRatio {
|
||||
/// The offending ratio.
|
||||
ratio: f64,
|
||||
},
|
||||
|
||||
/// A NaN or ±inf CSI value was found; corrupted input must be cleaned
|
||||
/// upstream, never masked over.
|
||||
#[error("Non-finite CSI value {value} at (t={row}, sc={col})")]
|
||||
NonFiniteValue {
|
||||
/// Time index of the offending value.
|
||||
row: usize,
|
||||
/// Subcarrier index of the offending value.
|
||||
col: usize,
|
||||
/// The non-finite value itself.
|
||||
value: f32,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -49,11 +49,13 @@ pub mod domain;
|
||||
pub mod error;
|
||||
pub mod eval;
|
||||
pub mod geometry;
|
||||
pub mod mae;
|
||||
pub mod rapid_adapt;
|
||||
pub mod ruview_metrics;
|
||||
pub mod signal_features;
|
||||
pub mod subcarrier;
|
||||
pub mod virtual_aug;
|
||||
pub mod wiflow_std;
|
||||
|
||||
// The following modules use `tch` (PyTorch Rust bindings) for GPU-accelerated
|
||||
// training and are only compiled when the `tch-backend` feature is enabled.
|
||||
@@ -81,7 +83,7 @@ pub use config::TrainingConfig;
|
||||
pub use dataset::{
|
||||
CsiDataset, CsiSample, DataLoader, MmFiDataset, SyntheticConfig, SyntheticCsiDataset,
|
||||
};
|
||||
pub use error::{ConfigError, DatasetError, SubcarrierError, TrainError};
|
||||
pub use error::{ConfigError, DatasetError, MaeError, SubcarrierError, TrainError};
|
||||
// TrainResult<T> is the generic Result alias from error.rs; the concrete
|
||||
// TrainResult struct from trainer.rs is accessed via trainer::TrainResult.
|
||||
pub use error::TrainResult as TrainResultAlias;
|
||||
@@ -89,6 +91,14 @@ pub use subcarrier::{
|
||||
compute_interp_weights, interpolate_subcarriers, select_subcarriers_by_variance,
|
||||
};
|
||||
|
||||
// ADR-152 §2.3 — UNSW MAE pretraining recipe re-exports.
|
||||
pub use mae::{patchify, random_mask, unpatchify, MaePretrainConfig, MaskIndices, PatchGrid};
|
||||
|
||||
// ADR-152 §2.2 — WiFlow-STD (DY2434) spatio-temporal-decoupled pose model.
|
||||
pub use wiflow_std::WiFlowStdConfig;
|
||||
#[cfg(feature = "tch-backend")]
|
||||
pub use wiflow_std::WiFlowStdModel;
|
||||
|
||||
// MERIDIAN (ADR-027) re-exports.
|
||||
pub use domain::{AdversarialSchedule, DomainClassifier, DomainFactorizer, GradientReversalLayer};
|
||||
pub use eval::CrossDomainEvaluator;
|
||||
|
||||
@@ -118,7 +118,7 @@ impl WiFiDensePoseLoss {
|
||||
// Normalise by number of visible joints in the batch.
|
||||
let n_visible = visibility.sum(Kind::Float);
|
||||
// Guard against division by zero (entire batch may have no labels).
|
||||
let safe_n = n_visible.clamp(1.0, f64::MAX);
|
||||
let safe_n = n_visible.clamp_min(1.0);
|
||||
|
||||
masked.sum(Kind::Float) / safe_n
|
||||
}
|
||||
@@ -165,7 +165,7 @@ impl WiFiDensePoseLoss {
|
||||
let masked_target_uv = target_uv * &fg_mask_f;
|
||||
|
||||
// Count foreground pixels × 48 channels to normalise.
|
||||
let n_fg = fg_mask_f.sum(Kind::Float).clamp(1.0, f64::MAX);
|
||||
let n_fg = fg_mask_f.sum(Kind::Float).clamp_min(1.0);
|
||||
|
||||
// Smooth-L1 with beta=1.0, reduction=Sum then divide by fg count.
|
||||
let uv_loss_sum = masked_pred_uv.smooth_l1_loss(&masked_target_uv, Reduction::Sum, 1.0);
|
||||
@@ -234,7 +234,7 @@ impl WiFiDensePoseLoss {
|
||||
// UV loss (foreground masked)
|
||||
let fg_mask = target_int.not_equal(0_i64);
|
||||
let fg_mask_f = fg_mask.unsqueeze(1).expand_as(pu).to_kind(Kind::Float);
|
||||
let n_fg = fg_mask_f.sum(Kind::Float).clamp(1.0, f64::MAX);
|
||||
let n_fg = fg_mask_f.sum(Kind::Float).clamp_min(1.0);
|
||||
let uv_loss =
|
||||
(pu * &fg_mask_f).smooth_l1_loss(&(tu * &fg_mask_f), Reduction::Sum, 1.0)
|
||||
/ n_fg;
|
||||
@@ -743,10 +743,11 @@ mod tests {
|
||||
}
|
||||
|
||||
// Visible batch (index 1) should have non-zero heatmaps.
|
||||
let heatmaps_ref = &heatmaps;
|
||||
let batch1_sum: f32 = (0..num_joints)
|
||||
.map(|j| {
|
||||
(0..size)
|
||||
.flat_map(|r| (0..size).map(move |c| heatmaps[[1, j, r, c]]))
|
||||
.flat_map(|r| (0..size).map(move |c| heatmaps_ref[[1, j, r, c]]))
|
||||
.sum::<f32>()
|
||||
})
|
||||
.sum();
|
||||
|
||||
@@ -0,0 +1,396 @@
|
||||
//! Masked-autoencoder (MAE) pretraining recipe for the ADR-150 RF foundation
|
||||
//! encoder — ADR-152 §2.3 (amends ADR-150 §2.3).
|
||||
//!
|
||||
//! Implements the *measured* tokenization recipe from the UNSW MAE pretraining
|
||||
//! study (arXiv [2511.18792](https://arxiv.org/abs/2511.18792), Nov 2025), the
|
||||
//! largest heterogeneous CSI pretraining run to date (1,320,892 samples, 14
|
||||
//! public datasets, 4 devices, 2.4/5/6 GHz, 20–160 MHz):
|
||||
//!
|
||||
//! - **80% masking ratio** over the patch grid.
|
||||
//! - **Small (30, 3) patches** — 30 time steps × 3 subcarriers — measured
|
||||
//! **+4.7%** over (40, 5) patches by preserving fine temporal dynamics.
|
||||
//! - Encoder capacity stays **ViT-Small-class (~15M params)**: ViT-Base adds
|
||||
//! only +0.4–0.9% over ViT-Small in-study, corroborating ADR-150's own
|
||||
//! finding that capacity hurts cross-subject transfer.
|
||||
//! - Unseen-domain performance scales **log-linearly with pretraining data,
|
||||
//! unsaturated at 1.3M samples** — data aggregation outranks architecture
|
||||
//! work (ADR-152 §2.3).
|
||||
//!
|
||||
//! This module provides the GPU-free half of the recipe: configuration,
|
||||
//! patchification, and deterministic random masking. The (future, ADR-150)
|
||||
//! encoder consumes [`PatchGrid`] + [`MaskIndices`] to compute the masked
|
||||
//! reconstruction loss (`L_masked_csi` in ADR-150 §2.3's loss stack).
|
||||
//!
|
||||
//! ## Axis convention
|
||||
//!
|
||||
//! A CSI window is `time × subcarriers`, row-major (`index = t * subc + sc`),
|
||||
//! matching the crate's `[T, …, n_sc]` dataset layout (time first, subcarriers
|
||||
//! last) and the UNSW "(30 time steps, 3 subcarriers)" patch framing. Patches
|
||||
//! are indexed row-major over the patch grid (`p = pt * n_patches_subc + ps`),
|
||||
//! and values within a patch are row-major time-major
|
||||
//! (`local = lt * patch_subc + lsc`).
|
||||
//!
|
||||
//! ## Divisibility policy: error, never truncate
|
||||
//!
|
||||
//! Window dimensions **must** be exact multiples of the patch dimensions.
|
||||
//! Non-divisible shapes return [`MaeError::NotDivisible`] instead of silently
|
||||
//! truncating trailing samples (this crate never silently drops data). The
|
||||
//! error names the largest divisible crop; use
|
||||
//! [`MaePretrainConfig::cropped_window_shape`] to compute it and crop
|
||||
//! explicitly before calling [`patchify`].
|
||||
//!
|
||||
//! ## Example
|
||||
//!
|
||||
//! ```rust
|
||||
//! use wifi_densepose_train::mae::MaePretrainConfig;
|
||||
//!
|
||||
//! let cfg = MaePretrainConfig::default(); // 0.80 masking, (30, 3) patches
|
||||
//! cfg.validate().expect("default recipe is valid");
|
||||
//!
|
||||
//! // 90 frames × 54 subcarriers → a 3 × 18 grid of (30, 3) patches.
|
||||
//! let window = vec![0.25_f32; 90 * 54];
|
||||
//! let (grid, mask) = cfg.mask_window(&window, 90, 54).unwrap();
|
||||
//! assert_eq!(grid.n_patches(), 54);
|
||||
//! assert_eq!(mask.masked.len(), 43); // round(0.80 * 54)
|
||||
//! assert_eq!(mask.visible.len(), 11);
|
||||
//! ```
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::error::{ConfigError, MaeError};
|
||||
use crate::virtual_aug::Xorshift64;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// MaePretrainConfig
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Hyper-parameters for masked-CSI pretraining (ADR-152 §2.3).
|
||||
///
|
||||
/// Defaults are the measured-optimal UNSW recipe (arXiv 2511.18792); change
|
||||
/// them only with benchmark evidence. Serializable so the recipe is recorded
|
||||
/// in checkpoint metadata alongside [`crate::config::TrainingConfig`].
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct MaePretrainConfig {
|
||||
/// Fraction of patches hidden from the encoder, in `(0, 1)`.
|
||||
///
|
||||
/// Default: **0.80** (UNSW measured optimum).
|
||||
pub mask_ratio: f64,
|
||||
|
||||
/// Patch extent along the time axis, in frames. Default: **30**.
|
||||
pub patch_time: usize,
|
||||
|
||||
/// Patch extent along the subcarrier axis. Default: **3**.
|
||||
pub patch_subc: usize,
|
||||
|
||||
/// Base seed for the deterministic mask sampler. Default: **42**.
|
||||
///
|
||||
/// For per-sample masks derive a child seed (e.g.
|
||||
/// `seed ^ sample_idx as u64`) and pass it to [`random_mask`]; reusing one
|
||||
/// seed yields the identical mask for every sample.
|
||||
pub seed: u64,
|
||||
}
|
||||
|
||||
impl Default for MaePretrainConfig {
|
||||
fn default() -> Self {
|
||||
MaePretrainConfig {
|
||||
mask_ratio: 0.80,
|
||||
patch_time: 30,
|
||||
patch_subc: 3,
|
||||
seed: 42,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl MaePretrainConfig {
|
||||
/// Validate the shape-independent fields.
|
||||
///
|
||||
/// # Validated invariants
|
||||
///
|
||||
/// - `mask_ratio` must be strictly inside `(0, 1)` and finite.
|
||||
/// - `patch_time` and `patch_subc` must be at least 1.
|
||||
pub fn validate(&self) -> Result<(), ConfigError> {
|
||||
if !self.mask_ratio.is_finite() || self.mask_ratio <= 0.0 || self.mask_ratio >= 1.0 {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"mask_ratio",
|
||||
format!("must be in (0.0, 1.0), got {}", self.mask_ratio),
|
||||
));
|
||||
}
|
||||
if self.patch_time == 0 {
|
||||
return Err(ConfigError::invalid_value("patch_time", "must be >= 1"));
|
||||
}
|
||||
if self.patch_subc == 0 {
|
||||
return Err(ConfigError::invalid_value("patch_subc", "must be >= 1"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check this recipe against a concrete `time × subc` window shape.
|
||||
///
|
||||
/// Errors if a patch dimension exceeds the window or if either axis is
|
||||
/// not an exact multiple of the patch extent (divisibility policy above).
|
||||
pub fn validate_for_window(&self, time: usize, subc: usize) -> Result<(), MaeError> {
|
||||
check_axis("time", time, self.patch_time)?;
|
||||
check_axis("subcarrier", subc, self.patch_subc)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Largest `(time, subc)` crop of the given window that is exactly
|
||||
/// divisible by the patch dimensions. Either component may be 0 when the
|
||||
/// window is smaller than one patch.
|
||||
#[must_use]
|
||||
pub fn cropped_window_shape(&self, time: usize, subc: usize) -> (usize, usize) {
|
||||
(
|
||||
(time / self.patch_time) * self.patch_time,
|
||||
(subc / self.patch_subc) * self.patch_subc,
|
||||
)
|
||||
}
|
||||
|
||||
/// Number of patches a `time × subc` window yields under this recipe.
|
||||
pub fn num_patches(&self, time: usize, subc: usize) -> Result<usize, MaeError> {
|
||||
self.validate_for_window(time, subc)?;
|
||||
Ok((time / self.patch_time) * (subc / self.patch_subc))
|
||||
}
|
||||
|
||||
/// Exact number of masked patches for a grid of `n_patches`:
|
||||
/// `round(mask_ratio * n_patches)`, clamped to `[0, n_patches]`.
|
||||
#[must_use]
|
||||
pub fn num_masked(&self, n_patches: usize) -> usize {
|
||||
((self.mask_ratio * n_patches as f64).round() as usize).min(n_patches)
|
||||
}
|
||||
|
||||
/// Patchify `window` and draw the deterministic random mask in one step,
|
||||
/// using `self.seed`. See [`patchify`] and [`random_mask`].
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Everything [`patchify`] rejects, plus [`MaeError::InvalidMaskRatio`]
|
||||
/// if `self.mask_ratio` is not finite or outside `(0, 1)` (the
|
||||
/// [`Self::validate`] rule) — a NaN ratio must never silently mask zero
|
||||
/// patches.
|
||||
pub fn mask_window(
|
||||
&self,
|
||||
window: &[f32],
|
||||
time: usize,
|
||||
subc: usize,
|
||||
) -> Result<(PatchGrid, MaskIndices), MaeError> {
|
||||
let grid = patchify(window, time, subc, self)?;
|
||||
let mask = random_mask(grid.n_patches(), self.mask_ratio, self.seed)?;
|
||||
Ok((grid, mask))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// PatchGrid / MaskIndices
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A CSI window decomposed into non-overlapping `patch_time × patch_subc`
|
||||
/// patches (see the module-level axis convention).
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PatchGrid {
|
||||
/// Patch extent along the time axis.
|
||||
pub patch_time: usize,
|
||||
/// Patch extent along the subcarrier axis.
|
||||
pub patch_subc: usize,
|
||||
/// Number of patch rows (`time / patch_time`).
|
||||
pub n_patches_time: usize,
|
||||
/// Number of patch columns (`subc / patch_subc`).
|
||||
pub n_patches_subc: usize,
|
||||
/// Flattened patches, row-major over the grid; each inner `Vec` is one
|
||||
/// patch of length `patch_time * patch_subc`, row-major time-major.
|
||||
pub patches: Vec<Vec<f32>>,
|
||||
}
|
||||
|
||||
impl PatchGrid {
|
||||
/// Total number of patches in the grid.
|
||||
#[must_use]
|
||||
pub fn n_patches(&self) -> usize {
|
||||
self.n_patches_time * self.n_patches_subc
|
||||
}
|
||||
|
||||
/// Number of scalar values per patch.
|
||||
#[must_use]
|
||||
pub fn patch_len(&self) -> usize {
|
||||
self.patch_time * self.patch_subc
|
||||
}
|
||||
|
||||
/// Window shape `(time, subc)` this grid reconstructs to.
|
||||
#[must_use]
|
||||
pub fn window_shape(&self) -> (usize, usize) {
|
||||
(
|
||||
self.n_patches_time * self.patch_time,
|
||||
self.n_patches_subc * self.patch_subc,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Sorted, disjoint patch-index sets produced by [`random_mask`]. Together
|
||||
/// they cover `0..n_patches` exactly.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct MaskIndices {
|
||||
/// Indices of patches hidden from the encoder (`round(ratio * n)` of them).
|
||||
pub masked: Vec<usize>,
|
||||
/// Indices of patches the encoder sees.
|
||||
pub visible: Vec<usize>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// patchify / unpatchify
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Decompose a row-major `time × subc` CSI window into the patch grid defined
|
||||
/// by `cfg`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// - [`MaeError::WindowShapeMismatch`] if `window.len() != time * subc`.
|
||||
/// - [`MaeError::PatchExceedsWindow`] / [`MaeError::NotDivisible`] per the
|
||||
/// module-level divisibility policy.
|
||||
/// - [`MaeError::NonFiniteValue`] on the first NaN/±inf encountered —
|
||||
/// corrupted CSI must be cleaned upstream, never masked over (cf. the
|
||||
/// WiFlow-STD NaN-poisoning incident, ADR-152 §2.2).
|
||||
pub fn patchify(
|
||||
window: &[f32],
|
||||
time: usize,
|
||||
subc: usize,
|
||||
cfg: &MaePretrainConfig,
|
||||
) -> Result<PatchGrid, MaeError> {
|
||||
let expected = time * subc;
|
||||
if window.len() != expected {
|
||||
return Err(MaeError::WindowShapeMismatch {
|
||||
time,
|
||||
subc,
|
||||
expected,
|
||||
actual: window.len(),
|
||||
});
|
||||
}
|
||||
cfg.validate_for_window(time, subc)?;
|
||||
if let Some(idx) = window.iter().position(|v| !v.is_finite()) {
|
||||
return Err(MaeError::NonFiniteValue {
|
||||
row: idx / subc,
|
||||
col: idx % subc,
|
||||
value: window[idx],
|
||||
});
|
||||
}
|
||||
|
||||
let n_patches_time = time / cfg.patch_time;
|
||||
let n_patches_subc = subc / cfg.patch_subc;
|
||||
let mut patches = Vec::with_capacity(n_patches_time * n_patches_subc);
|
||||
for pt in 0..n_patches_time {
|
||||
for ps in 0..n_patches_subc {
|
||||
let mut patch = Vec::with_capacity(cfg.patch_time * cfg.patch_subc);
|
||||
for lt in 0..cfg.patch_time {
|
||||
let t = pt * cfg.patch_time + lt;
|
||||
let row_start = t * subc + ps * cfg.patch_subc;
|
||||
patch.extend_from_slice(&window[row_start..row_start + cfg.patch_subc]);
|
||||
}
|
||||
patches.push(patch);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(PatchGrid {
|
||||
patch_time: cfg.patch_time,
|
||||
patch_subc: cfg.patch_subc,
|
||||
n_patches_time,
|
||||
n_patches_subc,
|
||||
patches,
|
||||
})
|
||||
}
|
||||
|
||||
/// Reassemble the full row-major `time × subc` window from a [`PatchGrid`].
|
||||
/// Exact inverse of [`patchify`].
|
||||
#[must_use]
|
||||
pub fn unpatchify(grid: &PatchGrid) -> Vec<f32> {
|
||||
unpatchify_select(grid, None, 0.0)
|
||||
}
|
||||
|
||||
/// Reassemble the window keeping only the patches listed in `visible`;
|
||||
/// every other patch's region is filled with `fill` (the standard MAE
|
||||
/// "visible tokens + mask token" view of the input).
|
||||
#[must_use]
|
||||
pub fn unpatchify_visible(grid: &PatchGrid, visible: &[usize], fill: f32) -> Vec<f32> {
|
||||
unpatchify_select(grid, Some(visible), fill)
|
||||
}
|
||||
|
||||
fn unpatchify_select(grid: &PatchGrid, keep: Option<&[usize]>, fill: f32) -> Vec<f32> {
|
||||
let (time, subc) = grid.window_shape();
|
||||
let mut window = vec![fill; time * subc];
|
||||
for (p, patch) in grid.patches.iter().enumerate() {
|
||||
if let Some(keep) = keep {
|
||||
if !keep.contains(&p) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
let pt = p / grid.n_patches_subc;
|
||||
let ps = p % grid.n_patches_subc;
|
||||
for lt in 0..grid.patch_time {
|
||||
let t = pt * grid.patch_time + lt;
|
||||
let row_start = t * subc + ps * grid.patch_subc;
|
||||
let local_start = lt * grid.patch_subc;
|
||||
window[row_start..row_start + grid.patch_subc]
|
||||
.copy_from_slice(&patch[local_start..local_start + grid.patch_subc]);
|
||||
}
|
||||
}
|
||||
window
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// random_mask
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Draw a deterministic random mask over `n_patches` patches.
|
||||
///
|
||||
/// Exactly `round(mask_ratio * n_patches)` patches (clamped to
|
||||
/// `[0, n_patches]`) are masked, chosen by a seeded Fisher–Yates shuffle
|
||||
/// ([`Xorshift64`]), so the same `(n_patches, mask_ratio, seed)` triple always
|
||||
/// yields the same mask. Both index lists are sorted ascending, disjoint, and
|
||||
/// together cover `0..n_patches`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// [`MaeError::InvalidMaskRatio`] if `mask_ratio` is not finite or outside
|
||||
/// the open interval `(0, 1)` — the same rule as
|
||||
/// [`MaePretrainConfig::validate`]. Erroring (never clamping) keeps the
|
||||
/// module's error-not-silent policy: a NaN ratio would otherwise silently
|
||||
/// mask zero patches and a ratio ≥ 1 would mask everything.
|
||||
pub fn random_mask(n_patches: usize, mask_ratio: f64, seed: u64) -> Result<MaskIndices, MaeError> {
|
||||
if !mask_ratio.is_finite() || mask_ratio <= 0.0 || mask_ratio >= 1.0 {
|
||||
return Err(MaeError::InvalidMaskRatio { ratio: mask_ratio });
|
||||
}
|
||||
let n_masked = ((mask_ratio * n_patches as f64).round() as usize).min(n_patches);
|
||||
let mut order: Vec<usize> = (0..n_patches).collect();
|
||||
let mut rng = Xorshift64::new(seed);
|
||||
for i in (1..n_patches).rev() {
|
||||
let j = (rng.next_u64() % (i as u64 + 1)) as usize;
|
||||
order.swap(i, j);
|
||||
}
|
||||
let mut masked: Vec<usize> = order[..n_masked].to_vec();
|
||||
let mut visible: Vec<usize> = order[n_masked..].to_vec();
|
||||
masked.sort_unstable();
|
||||
visible.sort_unstable();
|
||||
Ok(MaskIndices { masked, visible })
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn check_axis(axis: &'static str, window: usize, patch: usize) -> Result<(), MaeError> {
|
||||
if patch > window {
|
||||
return Err(MaeError::PatchExceedsWindow {
|
||||
axis,
|
||||
patch,
|
||||
window,
|
||||
});
|
||||
}
|
||||
let remainder = window % patch;
|
||||
if remainder != 0 {
|
||||
return Err(MaeError::NotDivisible {
|
||||
axis,
|
||||
window,
|
||||
patch,
|
||||
remainder,
|
||||
crop: window - remainder,
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -19,6 +19,7 @@
|
||||
|
||||
use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
|
||||
use petgraph::graph::{DiGraph, NodeIndex};
|
||||
use petgraph::visit::EdgeRef;
|
||||
use ruvector_mincut::{DynamicMinCut, MinCutBuilder};
|
||||
use std::collections::VecDeque;
|
||||
|
||||
@@ -106,6 +107,24 @@ impl Default for MetricsResult {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// EvalMetrics
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Per-evaluation pose metrics.
|
||||
///
|
||||
/// Plain value container produced by evaluation runs: lower `mpjpe`/`gps`
|
||||
/// and higher `pck_at_05` indicate better predictions.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq)]
|
||||
pub struct EvalMetrics {
|
||||
/// Mean Per-Joint Position Error (normalised units).
|
||||
pub mpjpe: f64,
|
||||
/// Percentage of Correct Keypoints at threshold 0.05 (0-1 scale).
|
||||
pub pck_at_05: f64,
|
||||
/// Geodesic Point Similarity error for DensePose surface predictions.
|
||||
pub gps: f64,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// MetricsAccumulator
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -126,7 +126,15 @@ impl WiFiDensePoseModel {
|
||||
tch::no_grad(|| self.forward_impl(amplitude, phase, false))
|
||||
}
|
||||
|
||||
/// Save model weights to a file (tch safetensors / .pt format).
|
||||
/// Save model weights to a file. The tch `VarStore` dispatches the format
|
||||
/// on the file extension: `.safetensors` → safetensors, anything else →
|
||||
/// torch `.pt`.
|
||||
///
|
||||
/// **Platform constraint:** prefer `.safetensors`. The `.pt` path
|
||||
/// (`_save_parameters`/`_load_parameters`) is broken on Windows with
|
||||
/// torch 2.11 (GenericDict internal assert on the load roundtrip — see
|
||||
/// `wiflow_std/model.rs::save_and_load_roundtrip`), which is why
|
||||
/// [`crate::trainer::Trainer`] writes `.safetensors` checkpoints.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
@@ -137,7 +145,8 @@ impl WiFiDensePoseModel {
|
||||
.map_err(|e| TrainError::training_step(format!("save failed: {e}")))
|
||||
}
|
||||
|
||||
/// Load model weights from a file.
|
||||
/// Load model weights from a file (format dispatched on extension; see
|
||||
/// the `.pt`-on-Windows caveat on [`Self::save`]).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
@@ -182,7 +191,7 @@ impl WiFiDensePoseModel {
|
||||
self.vs
|
||||
.trainable_variables()
|
||||
.iter()
|
||||
.map(|t| t.numel())
|
||||
.map(|t| t.numel() as i64)
|
||||
.sum()
|
||||
}
|
||||
|
||||
@@ -297,7 +306,12 @@ fn apply_antenna_attention(x: &Tensor, lambda: f32) -> Tensor {
|
||||
let xi = x.select(0, bi as i64); // [n_ant, n_sc]
|
||||
|
||||
// Move to CPU and convert to f32 for the pure-Rust attention kernel.
|
||||
let flat: Vec<f32> = Vec::from(xi.to_kind(Kind::Float).to_device(Device::Cpu).contiguous());
|
||||
let flat: Vec<f32> = Vec::<f32>::try_from(
|
||||
xi.to_kind(Kind::Float)
|
||||
.to_device(Device::Cpu)
|
||||
.flatten(0, -1),
|
||||
)
|
||||
.expect("antenna tensor to vec");
|
||||
|
||||
// Q = K = V = the antenna features (self-attention over antenna paths).
|
||||
let out = attn_mincut(
|
||||
@@ -350,7 +364,12 @@ fn apply_spatial_attention(x: &Tensor) -> Tensor {
|
||||
for bi in 0..b {
|
||||
// Extract [C, H*W] and transpose to [H*W, C].
|
||||
let xi = x.select(0, bi).reshape([c, h * w]).transpose(0, 1); // [H*W, C]
|
||||
let flat: Vec<f32> = Vec::from(xi.to_kind(Kind::Float).to_device(Device::Cpu).contiguous());
|
||||
let flat: Vec<f32> = Vec::<f32>::try_from(
|
||||
xi.to_kind(Kind::Float)
|
||||
.to_device(Device::Cpu)
|
||||
.flatten(0, -1),
|
||||
)
|
||||
.expect("spatial tensor to vec");
|
||||
|
||||
// Build token slices — one per spatial position.
|
||||
let tokens: Vec<&[f32]> = (0..n_spatial).map(|i| &flat[i * d..(i + 1) * d]).collect();
|
||||
@@ -973,7 +992,9 @@ mod tests {
|
||||
let mut model = WiFiDensePoseModel::new(&cfg, Device::Cpu);
|
||||
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let path = tmp.path().join("weights.pt");
|
||||
// safetensors, not .pt: this torch build's .pt roundtrip is broken on
|
||||
// Windows (torch 2.11 GenericDict internal assert).
|
||||
let path = tmp.path().join("weights.safetensors");
|
||||
|
||||
model.save(&path).expect("save should succeed");
|
||||
model.load(&path).expect("load should succeed");
|
||||
|
||||
@@ -153,11 +153,11 @@ pub fn run_proof(proof_dir: &Path) -> Result<ProofResult, Box<dyn std::error::Er
|
||||
let num_kp = kp.size()[1] as usize;
|
||||
let hm_size = cfg.heatmap_size;
|
||||
|
||||
let kp_vec: Vec<f32> = Vec::<f64>::from(kp.to_kind(Kind::Double).flatten(0, -1))
|
||||
let kp_vec: Vec<f32> = Vec::<f64>::try_from(kp.to_kind(Kind::Double).flatten(0, -1))?
|
||||
.iter()
|
||||
.map(|&x| x as f32)
|
||||
.collect();
|
||||
let vis_vec: Vec<f32> = Vec::<f64>::from(vis.to_kind(Kind::Double).flatten(0, -1))
|
||||
let vis_vec: Vec<f32> = Vec::<f64>::try_from(vis.to_kind(Kind::Double).flatten(0, -1))?
|
||||
.iter()
|
||||
.map(|&x| x as f32)
|
||||
.collect();
|
||||
@@ -261,7 +261,7 @@ pub fn hash_model_weights(model: &WiFiDensePoseModel) -> String {
|
||||
.flatten(0, -1)
|
||||
.to_kind(Kind::Float)
|
||||
.to_device(Device::Cpu);
|
||||
let values: Vec<f32> = Vec::<f32>::from(&flat);
|
||||
let values: Vec<f32> = Vec::<f32>::try_from(&flat).expect("param tensor to vec");
|
||||
let mut buf = vec![0u8; values.len() * 4];
|
||||
for (i, v) in values.iter().enumerate() {
|
||||
let bytes = v.to_le_bytes();
|
||||
@@ -292,6 +292,15 @@ pub fn load_expected_hash(proof_dir: &Path) -> Result<Option<String>, std::io::E
|
||||
Ok(if hash.is_empty() { None } else { Some(hash) })
|
||||
}
|
||||
|
||||
/// Verify that `path` is a valid checkpoint directory.
|
||||
///
|
||||
/// Returns `true` only when the path exists and is a directory. Deterministic
|
||||
/// and side-effect free — repeated calls always return the same result for an
|
||||
/// unchanged filesystem.
|
||||
pub fn verify_checkpoint_dir(path: &Path) -> bool {
|
||||
path.is_dir()
|
||||
}
|
||||
|
||||
/// Save the expected model hash to `<proof_dir>/expected_proof.sha256`.
|
||||
///
|
||||
/// Creates `proof_dir` if it does not already exist.
|
||||
|
||||
@@ -286,7 +286,12 @@ impl Trainer {
|
||||
best_epoch = epoch;
|
||||
patience_counter = 0;
|
||||
|
||||
let ckpt_name = format!("best_epoch{epoch:04}_pck{val_pck:.4}.pt");
|
||||
// .safetensors, not .pt: VarStore dispatches the format on
|
||||
// the extension, and this torch build's .pt
|
||||
// _save_parameters/_load_parameters roundtrip is broken on
|
||||
// Windows (torch 2.11 GenericDict internal assert — see
|
||||
// wiflow_std/model.rs save_and_load_roundtrip).
|
||||
let ckpt_name = format!("best_epoch{epoch:04}_pck{val_pck:.4}.safetensors");
|
||||
let ckpt_path = self.config.checkpoint_dir.join(&ckpt_name);
|
||||
|
||||
match self.model.save(&ckpt_path) {
|
||||
@@ -339,8 +344,8 @@ impl Trainer {
|
||||
}
|
||||
}
|
||||
|
||||
// Save final model regardless.
|
||||
let final_ckpt = self.config.checkpoint_dir.join("final.pt");
|
||||
// Save final model regardless (.safetensors — see checkpoint note above).
|
||||
let final_ckpt = self.config.checkpoint_dir.join("final.safetensors");
|
||||
if let Err(e) = self.model.save(&final_ckpt) {
|
||||
warn!("Failed to save final model: {e}");
|
||||
}
|
||||
@@ -413,7 +418,8 @@ impl Trainer {
|
||||
.load(path)
|
||||
.map_err(|e| TrainError::checkpoint(e.to_string(), path))?;
|
||||
|
||||
// Try to parse the epoch from the filename (e.g. "best_epoch0042_pck0.7842.pt").
|
||||
// Try to parse the epoch from the filename, extension-agnostic
|
||||
// (e.g. "best_epoch0042_pck0.7842.safetensors").
|
||||
let epoch = path
|
||||
.file_stem()
|
||||
.and_then(|s| s.to_str())
|
||||
@@ -582,11 +588,13 @@ fn kp_to_heatmap_tensor(
|
||||
let num_kp = kp_tensor.size()[1] as usize;
|
||||
|
||||
// Convert to ndarray for generate_target_heatmaps.
|
||||
let kp_vec: Vec<f32> = Vec::<f64>::from(kp_tensor.to_kind(Kind::Double).flatten(0, -1))
|
||||
let kp_vec: Vec<f32> = Vec::<f64>::try_from(kp_tensor.to_kind(Kind::Double).flatten(0, -1))
|
||||
.expect("kp tensor to vec")
|
||||
.iter()
|
||||
.map(|&x| x as f32)
|
||||
.collect();
|
||||
let vis_vec: Vec<f32> = Vec::<f64>::from(vis_tensor.to_kind(Kind::Double).flatten(0, -1))
|
||||
let vis_vec: Vec<f32> = Vec::<f64>::try_from(vis_tensor.to_kind(Kind::Double).flatten(0, -1))
|
||||
.expect("vis tensor to vec")
|
||||
.iter()
|
||||
.map(|&x| x as f32)
|
||||
.collect();
|
||||
@@ -622,8 +630,8 @@ fn heatmap_to_keypoints(heatmaps: &Tensor) -> Tensor {
|
||||
let arg = flat.argmax(-1, false);
|
||||
|
||||
// Decompose linear index into (row, col).
|
||||
let row = (&arg / w).to_kind(Kind::Float); // [B, 17]
|
||||
let col = (&arg % w).to_kind(Kind::Float); // [B, 17]
|
||||
let row = arg.divide_scalar_mode(w, "floor").to_kind(Kind::Float); // [B, 17]
|
||||
let col = arg.remainder(w).to_kind(Kind::Float); // [B, 17]
|
||||
|
||||
// Normalize to [0, 1]
|
||||
let x = col / (w - 1) as f64;
|
||||
@@ -639,7 +647,8 @@ fn heatmap_to_keypoints(heatmaps: &Tensor) -> Tensor {
|
||||
fn extract_kp_ndarray(kp_tensor: &Tensor, batch_idx: usize) -> Array2<f32> {
|
||||
let num_kp = kp_tensor.size()[1] as usize;
|
||||
let row = kp_tensor.select(0, batch_idx as i64);
|
||||
let data: Vec<f32> = Vec::<f64>::from(row.to_kind(Kind::Double).flatten(0, -1))
|
||||
let data: Vec<f32> = Vec::<f64>::try_from(row.to_kind(Kind::Double).flatten(0, -1))
|
||||
.expect("kp tensor to vec")
|
||||
.iter()
|
||||
.map(|&v| v as f32)
|
||||
.collect();
|
||||
@@ -652,7 +661,8 @@ fn extract_kp_ndarray(kp_tensor: &Tensor, batch_idx: usize) -> Array2<f32> {
|
||||
fn extract_vis_ndarray(vis_tensor: &Tensor, batch_idx: usize) -> Array1<f32> {
|
||||
let num_kp = vis_tensor.size()[1] as usize;
|
||||
let row = vis_tensor.select(0, batch_idx as i64);
|
||||
let data: Vec<f32> = Vec::<f64>::from(row.to_kind(Kind::Double))
|
||||
let data: Vec<f32> = Vec::<f64>::try_from(row.to_kind(Kind::Double))
|
||||
.expect("vis tensor to vec")
|
||||
.iter()
|
||||
.map(|&v| v as f32)
|
||||
.collect();
|
||||
|
||||
@@ -0,0 +1,899 @@
|
||||
//! Configuration and pure-Rust shape/parameter math for WiFlow-STD
|
||||
//! (ADR-152 §2.2). See the [module docs](crate::wiflow_std) for provenance.
|
||||
//!
|
||||
//! Everything here compiles without the `tch-backend` feature so the
|
||||
//! architecture's invariants (parameter count, output shapes, divisibility
|
||||
//! constraints) are unit-testable under `--no-default-features`. The
|
||||
//! 15-keypoint default must yield exactly **2,225,042** parameters — the
|
||||
//! count verified against the upstream reference (`RESULTS.md`).
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::error::ConfigError;
|
||||
|
||||
/// TCN kernel size — fixed at 3 in the reference architecture.
|
||||
pub const TCN_KERNEL: usize = 3;
|
||||
|
||||
/// Dropout used inside the 2-D conv blocks (`Dropout2d`). The reference
|
||||
/// hardcodes 0.3 in `convnet.py` (the model-level `dropout` argument is only
|
||||
/// 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
|
||||
}
|
||||
|
||||
fn default_min_feature_width() -> usize {
|
||||
15
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// WiFlowStdConfig
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Hyper-parameters for the WiFlow-STD pose model (ADR-152 §2.2).
|
||||
///
|
||||
/// Defaults reproduce the verified upstream architecture exactly (2,225,042
|
||||
/// parameters, 15 keypoints). For RuView's ESP32 17-keypoint eval set
|
||||
/// (ADR-152 §2.2(b)) use [`WiFlowStdConfig::for_keypoints`]`(17)` — the
|
||||
/// keypoint count only changes the final adaptive pooling, not the parameter
|
||||
/// count, so retrained 15-keypoint weights remain shape-compatible.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct WiFlowStdConfig {
|
||||
/// CSI input feature dimension (subcarriers × antenna paths flattened).
|
||||
/// Must be divisible by [`Self::tcn_groups`]. Default: **540**.
|
||||
pub subcarriers: usize,
|
||||
|
||||
/// Temporal window length in CSI frames. Default: **20**.
|
||||
pub window: usize,
|
||||
|
||||
/// Output channels of each TCN level (dilation doubles per level:
|
||||
/// 1, 2, 4, 8, …). Every entry must be divisible by [`Self::tcn_groups`].
|
||||
/// Default: **[540, 440, 340, 240]** — the `models/` code values, *not*
|
||||
/// upstream `config.py`'s stale `[480, 360, 240]`.
|
||||
pub tcn_channels: Vec<usize>,
|
||||
|
||||
/// 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**. 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]**.
|
||||
pub conv_channels: Vec<usize>,
|
||||
|
||||
/// Attention head groups for the dual axial attention. Must divide the
|
||||
/// last entry of [`Self::conv_channels`]. Default: **8**.
|
||||
pub attention_groups: usize,
|
||||
|
||||
/// Number of 2-D keypoints produced. Default: **15** (upstream skeleton);
|
||||
/// use **17** for RuView's COCO-skeleton ESP32 eval set. Only changes the
|
||||
/// parameter-free final adaptive pool — never the trunk: the stride
|
||||
/// schedule is governed by [`Self::min_feature_width`], so 15- and
|
||||
/// 17-keypoint variants share the identical conv graph and weights
|
||||
/// (matching the validated Python protocol,
|
||||
/// `benchmarks/wiflow-std/remote/measb/train_measb.py`, which swaps only
|
||||
/// `avg_pool` and loads the pretrained state_dict `strict=True`).
|
||||
pub keypoints: usize,
|
||||
|
||||
/// Floor for the conv encoder's width downsampling: each
|
||||
/// `AsymmetricConvBlock` halves the width only while the result stays
|
||||
/// ≥ this value (see [`Self::conv_strides`]).
|
||||
///
|
||||
/// Default: **15** — the upstream constant. Provenance: the reference's
|
||||
/// four hardcoded stride-2 blocks exist because its 240-channel TCN
|
||||
/// output halves cleanly four times, 240 / 2⁴ = 15. The compact presets'
|
||||
/// schedules were derived with this same floor. Override only when
|
||||
/// designing a new trunk; do **not** couple it to [`Self::keypoints`] —
|
||||
/// the adaptive pool maps the decoder height to any keypoint count.
|
||||
#[serde(default = "default_min_feature_width")]
|
||||
pub min_feature_width: usize,
|
||||
|
||||
/// Elementwise dropout probability inside the TCN blocks, in `[0, 1)`.
|
||||
/// Default: **0.5** (the value used by our verified retraining run).
|
||||
pub dropout: f64,
|
||||
}
|
||||
|
||||
impl Default for WiFlowStdConfig {
|
||||
fn default() -> Self {
|
||||
WiFlowStdConfig {
|
||||
subcarriers: 540,
|
||||
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,
|
||||
min_feature_width: 15,
|
||||
dropout: 0.5,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl WiFlowStdConfig {
|
||||
/// Default architecture with a different keypoint count (e.g. 17 for the
|
||||
/// ESP32 COCO-skeleton eval set, ADR-152 §2.2(b)).
|
||||
///
|
||||
/// The trunk is untouched: [`Self::min_feature_width`] stays at the
|
||||
/// upstream floor of 15, so e.g. `for_keypoints(17)` keeps the trained
|
||||
/// `[2, 2, 2, 2]` stride schedule (feature width 15) and the adaptive
|
||||
/// pool maps 15 → 17 — exactly the validated Python protocol
|
||||
/// (`benchmarks/wiflow-std/remote/measb/train_measb.py`).
|
||||
pub fn for_keypoints(keypoints: usize) -> Self {
|
||||
WiFlowStdConfig {
|
||||
keypoints,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// **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
|
||||
///
|
||||
/// Returns [`ConfigError::InvalidValue`] naming the offending field.
|
||||
pub fn validate(&self) -> Result<(), ConfigError> {
|
||||
if self.subcarriers == 0 {
|
||||
return Err(ConfigError::invalid_value("subcarriers", "must be >= 1"));
|
||||
}
|
||||
if self.window == 0 {
|
||||
return Err(ConfigError::invalid_value("window", "must be >= 1"));
|
||||
}
|
||||
if self.tcn_groups == 0 {
|
||||
return Err(ConfigError::invalid_value("tcn_groups", "must be >= 1"));
|
||||
}
|
||||
// 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!(
|
||||
"{} is not divisible by tcn_groups={} (grouped conv requirement)",
|
||||
self.subcarriers, self.tcn_groups
|
||||
),
|
||||
));
|
||||
}
|
||||
if self.tcn_channels.is_empty() {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"tcn_channels",
|
||||
"must contain at least one level",
|
||||
));
|
||||
}
|
||||
for (i, &c) in self.tcn_channels.iter().enumerate() {
|
||||
if c == 0 || (fixed && c % self.tcn_groups != 0) {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"tcn_channels",
|
||||
format!(
|
||||
"level {i} has {c} channels; must be > 0 and divisible by tcn_groups={}",
|
||||
self.tcn_groups
|
||||
),
|
||||
));
|
||||
}
|
||||
}
|
||||
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",
|
||||
"must contain at least one block",
|
||||
));
|
||||
}
|
||||
if self.conv_channels.iter().any(|&c| c == 0) {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"conv_channels",
|
||||
"all blocks must have > 0 channels",
|
||||
));
|
||||
}
|
||||
let c_last = *self.conv_channels.last().expect("non-empty checked above");
|
||||
if self.attention_groups == 0 || c_last % self.attention_groups != 0 {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"attention_groups",
|
||||
format!(
|
||||
"{} must be >= 1 and divide the last conv channel count {c_last}",
|
||||
self.attention_groups
|
||||
),
|
||||
));
|
||||
}
|
||||
if c_last < 2 || c_last % 2 != 0 {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"conv_channels",
|
||||
format!("last block has {c_last} channels; decoder needs an even count >= 2"),
|
||||
));
|
||||
}
|
||||
if self.keypoints == 0 {
|
||||
return Err(ConfigError::invalid_value("keypoints", "must be >= 1"));
|
||||
}
|
||||
if self.min_feature_width == 0 {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"min_feature_width",
|
||||
"must be >= 1",
|
||||
));
|
||||
}
|
||||
if !self.dropout.is_finite() || !(0.0..1.0).contains(&self.dropout) {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"dropout",
|
||||
format!("{} is outside [0, 1)", self.dropout),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Shape inference
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Channel count produced by the TCN stack (last TCN level). This is the
|
||||
/// *width* of the image-like tensor fed to the 2-D encoder.
|
||||
pub fn tcn_output_channels(&self) -> usize {
|
||||
*self.tcn_channels.last().unwrap_or(&0)
|
||||
}
|
||||
|
||||
/// 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::min_feature_width`]. At the upstream
|
||||
/// default (240 TCN channels, floor 15) this derives `[2, 2, 2, 2]` —
|
||||
/// the hardcoded upstream schedule, exactly.
|
||||
///
|
||||
/// Deliberately independent of [`Self::keypoints`]: the keypoint count
|
||||
/// only changes the parameter-free adaptive pool, so retargeting the
|
||||
/// skeleton (e.g. [`Self::for_keypoints`]`(17)`) keeps the trained graph
|
||||
/// and the pool maps `feature_width() → keypoints`.
|
||||
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.min_feature_width {
|
||||
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 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 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).
|
||||
pub fn output_shape(&self, batch: usize) -> (usize, usize, usize) {
|
||||
(batch, self.keypoints, 2)
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Parameter-count formula
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Total trainable parameter count, derived layer-by-layer from the
|
||||
/// architecture (BatchNorm weight+bias counted; running stats are buffers
|
||||
/// and excluded, matching PyTorch's `numel` convention).
|
||||
///
|
||||
/// Pins the port against the verified reference: the 15-keypoint default
|
||||
/// must equal **2,225,042** (`RESULTS.md` artifact verification).
|
||||
///
|
||||
/// Returns **0** for any config that fails [`Self::validate`]: the
|
||||
/// formula is only meaningful for buildable architectures (an invalid
|
||||
/// config would otherwise index an empty `conv_channels` or divide by a
|
||||
/// zero group count). Call `validate()` first when you need the reason.
|
||||
pub fn param_count(&self) -> usize {
|
||||
if self.validate().is_err() {
|
||||
return 0;
|
||||
}
|
||||
|
||||
let mut total = 0;
|
||||
|
||||
// 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 (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;
|
||||
}
|
||||
|
||||
// ConvBlock1 (1 → conv_channels[0]) + asymmetric blocks. Both block
|
||||
// kinds have identical parameter shapes (stride changes nothing).
|
||||
let mut c_in = 1;
|
||||
total += conv_block_params(c_in, self.conv_channels[0]);
|
||||
c_in = self.conv_channels[0];
|
||||
for &c_out in &self.conv_channels {
|
||||
total += conv_block_params(c_in, c_out);
|
||||
c_in = c_out;
|
||||
}
|
||||
|
||||
// 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 → decoder_mid) + BN + 1×1 conv (mid → 2) + BN.
|
||||
total += decoder_params(c_in, self.decoder_mid());
|
||||
|
||||
total
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Per-component parameter formulas
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// 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. `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 / pw_groups); // pointwise 1×1
|
||||
let bn1p = 2 * c_out;
|
||||
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 / pw_groups) * c_out + 2 * c_out
|
||||
} else {
|
||||
0
|
||||
};
|
||||
grouped1 + bn1g + pw1 + bn1p + grouped2 + bn2g + pw2 + bn2p + downsample
|
||||
}
|
||||
|
||||
/// One `ConvBlock1` / `AsymmetricConvBlock`: three (1, 3) convs **with bias**
|
||||
/// + BN each, plus a bias-free 1×1 + BN residual projection.
|
||||
fn conv_block_params(c_in: usize, c_out: usize) -> usize {
|
||||
let conv1 = c_out * c_in * 3 + c_out;
|
||||
let conv_rest = 2 * (c_out * c_out * 3 + c_out);
|
||||
let bns = 3 * 2 * c_out;
|
||||
let downsample = c_in * c_out + 2 * c_out;
|
||||
conv1 + conv_rest + bns + downsample
|
||||
}
|
||||
|
||||
/// One `AxialAttention` axis: bias-free 1×1 qkv conv (c → 3c), BN over the
|
||||
/// 3c qkv channels, BN over the `groups` similarity maps, BN over the output.
|
||||
fn axial_attention_params(c: usize, groups: usize) -> usize {
|
||||
let qkv = c * 3 * c;
|
||||
let bn_qkv = 2 * (3 * c);
|
||||
let bn_similarity = 2 * groups;
|
||||
let bn_output = 2 * c;
|
||||
qkv + bn_qkv + bn_similarity + bn_output
|
||||
}
|
||||
|
||||
/// 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;
|
||||
let bn2 = 2 * 2;
|
||||
conv1 + bn1 + conv2 + bn2
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests (pure Rust — run under --no-default-features)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Reference parameter count verified against the upstream checkpoint
|
||||
/// and `torchinfo` (benchmarks/wiflow-std/RESULTS.md, 2026-06-10).
|
||||
const REFERENCE_PARAMS: usize = 2_225_042;
|
||||
|
||||
#[test]
|
||||
fn default_config_is_valid() {
|
||||
WiFlowStdConfig::default()
|
||||
.validate()
|
||||
.expect("default config must validate");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_param_count_matches_verified_reference() {
|
||||
assert_eq!(WiFlowStdConfig::default().param_count(), REFERENCE_PARAMS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn param_count_is_independent_of_keypoints() {
|
||||
// The keypoint count only changes the parameter-free adaptive pool,
|
||||
// so 15- and 17-keypoint variants share identical weights.
|
||||
let kp17 = WiFlowStdConfig::for_keypoints(17);
|
||||
kp17.validate().expect("17-keypoint config must validate");
|
||||
assert_eq!(kp17.param_count(), REFERENCE_PARAMS);
|
||||
}
|
||||
|
||||
#[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, 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);
|
||||
assert_eq!(conv_block_params(8, 16), 2_224);
|
||||
assert_eq!(conv_block_params(16, 32), 8_544);
|
||||
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, 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 for_keypoints_17_keeps_trained_trunk_and_pools_15_to_17() {
|
||||
// Pin against the validated Python protocol (train_measb.py): K=17
|
||||
// swaps only the adaptive pool, never the stride schedule. A derived
|
||||
// [2, 2, 2, 1]/width-30 graph here would silently diverge from the
|
||||
// trained [2, 2, 2, 2]/width-15 checkpoint.
|
||||
let cfg = WiFlowStdConfig::for_keypoints(17);
|
||||
assert_eq!(cfg.min_feature_width, 15);
|
||||
assert_eq!(cfg.conv_strides(), [2, 2, 2, 2]);
|
||||
assert_eq!(cfg.feature_width(), 15);
|
||||
assert_eq!(cfg.output_shape(1), (1, 17, 2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn min_feature_width_override_changes_schedule_as_designed() {
|
||||
// Raising the floor stops the downsampling earlier (240 → 30).
|
||||
let cfg = WiFlowStdConfig {
|
||||
min_feature_width: 30,
|
||||
..Default::default()
|
||||
};
|
||||
cfg.validate().expect("floor 30 validates");
|
||||
assert_eq!(cfg.conv_strides(), [2, 2, 2, 1]);
|
||||
assert_eq!(cfg.feature_width(), 30);
|
||||
|
||||
// Lowering it lets a small trunk halve further (tiny: 32 → 8).
|
||||
let cfg = WiFlowStdConfig {
|
||||
min_feature_width: 8,
|
||||
..WiFlowStdConfig::tiny()
|
||||
};
|
||||
cfg.validate().expect("floor 8 validates");
|
||||
assert_eq!(cfg.conv_strides(), [2, 2, 1, 1]);
|
||||
assert_eq!(cfg.feature_width(), 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_min_feature_width() {
|
||||
let cfg = WiFlowStdConfig {
|
||||
min_feature_width: 0,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(cfg.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn param_count_returns_zero_for_invalid_configs() {
|
||||
// Documented total behavior: configs that fail validate() yield 0
|
||||
// instead of panicking (OOB index / division by zero).
|
||||
for cfg in [
|
||||
WiFlowStdConfig {
|
||||
conv_channels: vec![],
|
||||
..Default::default()
|
||||
},
|
||||
WiFlowStdConfig {
|
||||
tcn_groups: 0,
|
||||
..Default::default()
|
||||
},
|
||||
WiFlowStdConfig {
|
||||
input_pw_groups: 0,
|
||||
..Default::default()
|
||||
},
|
||||
WiFlowStdConfig {
|
||||
tcn_channels: vec![],
|
||||
..Default::default()
|
||||
},
|
||||
] {
|
||||
assert!(cfg.validate().is_err(), "precondition: {cfg:?} is invalid");
|
||||
assert_eq!(cfg.param_count(), 0, "no panic, returns 0: {cfg:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[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]
|
||||
fn output_shape_default_and_esp32() {
|
||||
assert_eq!(WiFlowStdConfig::default().output_shape(4), (4, 15, 2));
|
||||
assert_eq!(
|
||||
WiFlowStdConfig::for_keypoints(17).output_shape(1),
|
||||
(1, 17, 2)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feature_width_default_is_15() {
|
||||
// 240 → 120 → 60 → 30 → 15 (four stride-(1,2) blocks).
|
||||
assert_eq!(WiFlowStdConfig::default().feature_width(), 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tcn_output_channels_default_is_240() {
|
||||
assert_eq!(WiFlowStdConfig::default().tcn_output_channels(), 240);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_subcarriers_not_divisible_by_groups() {
|
||||
let cfg = WiFlowStdConfig {
|
||||
subcarriers: 541,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(cfg.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_dimensions() {
|
||||
for cfg in [
|
||||
WiFlowStdConfig {
|
||||
subcarriers: 0,
|
||||
..Default::default()
|
||||
},
|
||||
WiFlowStdConfig {
|
||||
window: 0,
|
||||
..Default::default()
|
||||
},
|
||||
WiFlowStdConfig {
|
||||
keypoints: 0,
|
||||
..Default::default()
|
||||
},
|
||||
WiFlowStdConfig {
|
||||
tcn_groups: 0,
|
||||
..Default::default()
|
||||
},
|
||||
] {
|
||||
assert!(cfg.validate().is_err(), "expected rejection: {cfg:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_or_indivisible_tcn_channels() {
|
||||
let empty = WiFlowStdConfig {
|
||||
tcn_channels: vec![],
|
||||
..Default::default()
|
||||
};
|
||||
assert!(empty.validate().is_err());
|
||||
|
||||
let indivisible = WiFlowStdConfig {
|
||||
tcn_channels: vec![540, 441],
|
||||
..Default::default()
|
||||
};
|
||||
assert!(indivisible.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_bad_conv_channels() {
|
||||
let empty = WiFlowStdConfig {
|
||||
conv_channels: vec![],
|
||||
..Default::default()
|
||||
};
|
||||
assert!(empty.validate().is_err());
|
||||
|
||||
let zero = WiFlowStdConfig {
|
||||
conv_channels: vec![8, 0, 64],
|
||||
..Default::default()
|
||||
};
|
||||
assert!(zero.validate().is_err());
|
||||
|
||||
// Odd last channel breaks the c → c/2 decoder split.
|
||||
let odd_last = WiFlowStdConfig {
|
||||
conv_channels: vec![8, 16, 33],
|
||||
attention_groups: 1,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(odd_last.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_attention_group_mismatch() {
|
||||
let cfg = WiFlowStdConfig {
|
||||
attention_groups: 7, // 64 % 7 != 0
|
||||
..Default::default()
|
||||
};
|
||||
assert!(cfg.validate().is_err());
|
||||
let zero = WiFlowStdConfig {
|
||||
attention_groups: 0,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(zero.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_out_of_range_dropout() {
|
||||
for d in [1.0, 1.5, -0.1, f64::NAN] {
|
||||
let cfg = WiFlowStdConfig {
|
||||
dropout: d,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(cfg.validate().is_err(), "dropout {d} must be rejected");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serde_roundtrip_preserves_config() {
|
||||
let cfg = WiFlowStdConfig::for_keypoints(17);
|
||||
let json = serde_json::to_string(&cfg).expect("serialize");
|
||||
let back: WiFlowStdConfig = serde_json::from_str(&json).expect("deserialize");
|
||||
assert_eq!(back, cfg);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,334 @@
|
||||
//! Building-block layers for the WiFlow-STD model (tch backend, ADR-152 §2.2):
|
||||
//! grouped causal TCN blocks, asymmetric residual conv blocks, and dual axial
|
||||
//! attention. Internal to [`super::model`]; see the module docs for provenance.
|
||||
|
||||
use tch::{nn, nn::Module, Tensor};
|
||||
|
||||
use super::config::{CONV_BLOCK_DROPOUT, TCN_KERNEL};
|
||||
|
||||
/// BatchNorm config matching the reference: gamma = 1 (PyTorch default; the
|
||||
/// reference additionally pins BatchNorm1d weight=1/bias=0). tch-0.24's
|
||||
/// `BatchNormConfig::default()` would draw gamma from Uniform(0,1), silently
|
||||
/// halving activations on average in from-scratch training.
|
||||
pub(super) fn bn_cfg() -> nn::BatchNormConfig {
|
||||
nn::BatchNormConfig {
|
||||
ws_init: nn::Init::Const(1.0),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// GroupedTemporalBlock (TCN level)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// One TCN level: two (depthwise-grouped causal conv → BN → SiLU → pointwise
|
||||
/// conv → BN → SiLU → dropout) stages with a residual connection (1×1 + BN
|
||||
/// projection when channels change) and a final SiLU.
|
||||
///
|
||||
/// Causality: each grouped conv pads by `(k-1)·dilation` and the trailing
|
||||
/// padding is chomped off afterwards, exactly like the reference `Chomp1d`.
|
||||
pub(super) struct GroupedTemporalBlock {
|
||||
conv1_group: nn::Conv1D,
|
||||
bn1_group: nn::BatchNorm,
|
||||
conv1_pw: nn::Conv1D,
|
||||
bn1_pw: nn::BatchNorm,
|
||||
conv2_group: nn::Conv1D,
|
||||
bn2_group: nn::BatchNorm,
|
||||
conv2_pw: nn::Conv1D,
|
||||
bn2_pw: nn::BatchNorm,
|
||||
downsample: Option<(nn::Conv1D, nn::BatchNorm)>,
|
||||
dropout: f64,
|
||||
}
|
||||
|
||||
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,
|
||||
g_in: i64,
|
||||
g_out: i64,
|
||||
pw_groups: i64,
|
||||
dropout: f64,
|
||||
) -> Self {
|
||||
let k = TCN_KERNEL as i64;
|
||||
let padding = (k - 1) * dilation;
|
||||
let grouped_cfg = |groups| nn::ConvConfig {
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
bias: false,
|
||||
..Default::default()
|
||||
};
|
||||
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(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(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(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(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(pw_groups)),
|
||||
nn::batch_norm1d(&vs / "ds_bn", c_out, bn_cfg()),
|
||||
)
|
||||
});
|
||||
|
||||
GroupedTemporalBlock {
|
||||
conv1_group,
|
||||
bn1_group,
|
||||
conv1_pw,
|
||||
bn1_pw,
|
||||
conv2_group,
|
||||
bn2_group,
|
||||
conv2_pw,
|
||||
bn2_pw,
|
||||
downsample,
|
||||
dropout,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn forward_t(&self, x: &Tensor, train: bool) -> Tensor {
|
||||
let res = match &self.downsample {
|
||||
Some((conv, bn)) => conv.forward(x).apply_t(bn, train),
|
||||
None => x.shallow_clone(),
|
||||
};
|
||||
let t = x.size()[2];
|
||||
|
||||
// Stage 1: grouped causal conv (chomp trailing padding) + pointwise.
|
||||
let out = self
|
||||
.conv1_group
|
||||
.forward(x)
|
||||
.narrow(2, 0, t) // Chomp1d
|
||||
.apply_t(&self.bn1_group, train)
|
||||
.silu()
|
||||
.apply(&self.conv1_pw)
|
||||
.apply_t(&self.bn1_pw, train)
|
||||
.silu()
|
||||
.dropout(self.dropout, train);
|
||||
|
||||
// Stage 2.
|
||||
let out = self
|
||||
.conv2_group
|
||||
.forward(&out)
|
||||
.narrow(2, 0, t) // Chomp1d
|
||||
.apply_t(&self.bn2_group, train)
|
||||
.silu()
|
||||
.apply(&self.conv2_pw)
|
||||
.apply_t(&self.bn2_pw, train)
|
||||
.silu()
|
||||
.dropout(self.dropout, train);
|
||||
|
||||
(out + res).silu()
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ConvBlock (ConvBlock1 / AsymmetricConvBlock)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Asymmetric residual conv block: three `(1, 3)` convs (only the subcarrier
|
||||
/// axis is convolved) with BN, SiLU and channel dropout, plus a 1×1 + BN
|
||||
/// residual projection. `stride_w == 1` reproduces the reference `ConvBlock1`,
|
||||
/// `stride_w == 2` the downsampling `AsymmetricConvBlock`.
|
||||
pub(super) struct ConvBlock {
|
||||
conv1: nn::Conv2D,
|
||||
bn1: nn::BatchNorm,
|
||||
conv2: nn::Conv2D,
|
||||
bn2: nn::BatchNorm,
|
||||
conv3: nn::Conv2D,
|
||||
bn3: nn::BatchNorm,
|
||||
ds_conv: nn::Conv2D,
|
||||
ds_bn: nn::BatchNorm,
|
||||
}
|
||||
|
||||
impl ConvBlock {
|
||||
pub(super) fn new(vs: nn::Path, c_in: i64, c_out: i64, stride_w: i64) -> Self {
|
||||
let asym = |stride_w| nn::ConvConfigND::<[i64; 2]> {
|
||||
stride: [1, stride_w],
|
||||
padding: [0, 1],
|
||||
..Default::default()
|
||||
};
|
||||
let conv1 = nn::conv(&vs / "conv1", c_in, c_out, [1, 3], asym(stride_w));
|
||||
let bn1 = nn::batch_norm2d(&vs / "bn1", c_out, bn_cfg());
|
||||
let conv2 = nn::conv(&vs / "conv2", c_out, c_out, [1, 3], asym(1));
|
||||
let bn2 = nn::batch_norm2d(&vs / "bn2", c_out, bn_cfg());
|
||||
let conv3 = nn::conv(&vs / "conv3", c_out, c_out, [1, 3], asym(1));
|
||||
let bn3 = nn::batch_norm2d(&vs / "bn3", c_out, bn_cfg());
|
||||
|
||||
let ds_conv = nn::conv(
|
||||
&vs / "ds_conv",
|
||||
c_in,
|
||||
c_out,
|
||||
[1, 1],
|
||||
nn::ConvConfigND::<[i64; 2]> {
|
||||
stride: [1, stride_w],
|
||||
bias: false,
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let ds_bn = nn::batch_norm2d(&vs / "ds_bn", c_out, bn_cfg());
|
||||
|
||||
ConvBlock {
|
||||
conv1,
|
||||
bn1,
|
||||
conv2,
|
||||
bn2,
|
||||
conv3,
|
||||
bn3,
|
||||
ds_conv,
|
||||
ds_bn,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn forward_t(&self, x: &Tensor, train: bool) -> Tensor {
|
||||
let identity = self.ds_conv.forward(x).apply_t(&self.ds_bn, train);
|
||||
let out = x
|
||||
.apply(&self.conv1)
|
||||
.apply_t(&self.bn1, train)
|
||||
.silu()
|
||||
.feature_dropout(CONV_BLOCK_DROPOUT, train) // Dropout2d
|
||||
.apply(&self.conv2)
|
||||
.apply_t(&self.bn2, train)
|
||||
.silu()
|
||||
.feature_dropout(CONV_BLOCK_DROPOUT, train)
|
||||
.apply(&self.conv3)
|
||||
.apply_t(&self.bn3, train);
|
||||
(out + identity).silu()
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Axial attention
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Single-axis self-attention with BN-normalised qkv, BN-normalised
|
||||
/// similarity logits and BN-normalised output. `width == true` attends along
|
||||
/// the last (W) axis, otherwise along the H axis; the other spatial axis is
|
||||
/// folded into the batch.
|
||||
pub(super) struct AxialAttention {
|
||||
qkv: nn::Conv1D,
|
||||
bn_qkv: nn::BatchNorm,
|
||||
bn_similarity: nn::BatchNorm,
|
||||
bn_output: nn::BatchNorm,
|
||||
out_planes: i64,
|
||||
groups: i64,
|
||||
width: bool,
|
||||
}
|
||||
|
||||
impl AxialAttention {
|
||||
pub(super) fn new(vs: nn::Path, planes: i64, groups: i64, width: bool) -> Self {
|
||||
// Reference init: N(0, sqrt(1 / in_planes)).
|
||||
let qkv = nn::conv1d(
|
||||
&vs / "qkv",
|
||||
planes,
|
||||
planes * 3,
|
||||
1,
|
||||
nn::ConvConfig {
|
||||
bias: false,
|
||||
ws_init: nn::Init::Randn {
|
||||
mean: 0.0,
|
||||
stdev: (1.0 / planes as f64).sqrt(),
|
||||
},
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let bn_qkv = nn::batch_norm1d(&vs / "bn_qkv", planes * 3, bn_cfg());
|
||||
let bn_similarity = nn::batch_norm2d(&vs / "bn_similarity", groups, bn_cfg());
|
||||
let bn_output = nn::batch_norm1d(&vs / "bn_output", planes, bn_cfg());
|
||||
|
||||
AxialAttention {
|
||||
qkv,
|
||||
bn_qkv,
|
||||
bn_similarity,
|
||||
bn_output,
|
||||
out_planes: planes,
|
||||
groups,
|
||||
width,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn forward_t(&self, x: &Tensor, train: bool) -> Tensor {
|
||||
// Fold the non-attended spatial axis into the batch:
|
||||
// width: [B,C,H,W] → [B,H,C,W]; height: [B,C,H,W] → [B,W,C,H].
|
||||
let x = if self.width {
|
||||
x.permute([0, 2, 1, 3])
|
||||
} else {
|
||||
x.permute([0, 3, 1, 2])
|
||||
};
|
||||
let (n, outer, c, axis) = {
|
||||
let s = x.size();
|
||||
(s[0], s[1], s[2], s[3])
|
||||
};
|
||||
let flat = x.contiguous().view([n * outer, c, axis]);
|
||||
|
||||
// BN-normalised qkv: [N', 3·C, axis] → grouped q, k, v.
|
||||
let gp = self.out_planes / self.groups; // group planes
|
||||
let qkv = flat.apply(&self.qkv).apply_t(&self.bn_qkv, train).reshape([
|
||||
n * outer,
|
||||
3,
|
||||
self.groups,
|
||||
gp,
|
||||
axis,
|
||||
]);
|
||||
let q = qkv.select(1, 0); // [N', g, gp, axis]
|
||||
let k = qkv.select(1, 1);
|
||||
let v = qkv.select(1, 2);
|
||||
|
||||
// similarity[b,g,i,j] = Σ_c q[b,g,c,i]·k[b,g,c,j], BN over the g maps.
|
||||
let logits = q.transpose(2, 3).matmul(&k); // [N', g, axis, axis]
|
||||
let similarity = logits
|
||||
.apply_t(&self.bn_similarity, train)
|
||||
.softmax(-1, logits.kind());
|
||||
|
||||
// out[b,g,c,i] = Σ_j similarity[b,g,i,j]·v[b,g,c,j].
|
||||
let sv = v.matmul(&similarity.transpose(2, 3)); // [N', g, gp, axis]
|
||||
let out = sv
|
||||
.reshape([n * outer, self.out_planes, axis])
|
||||
.apply_t(&self.bn_output, train)
|
||||
.view([n, outer, self.out_planes, axis]);
|
||||
|
||||
// Restore [B, C, H, W].
|
||||
if self.width {
|
||||
out.permute([0, 2, 1, 3])
|
||||
} else {
|
||||
out.permute([0, 2, 3, 1])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Width-axis then height-axis axial attention (the reference
|
||||
/// `DualAxialAttention`, stride 1).
|
||||
pub(super) struct DualAxialAttention {
|
||||
width_axis: AxialAttention,
|
||||
height_axis: AxialAttention,
|
||||
}
|
||||
|
||||
impl DualAxialAttention {
|
||||
pub(super) fn new(vs: nn::Path, planes: i64, groups: i64) -> Self {
|
||||
DualAxialAttention {
|
||||
width_axis: AxialAttention::new(&vs / "width", planes, groups, true),
|
||||
height_axis: AxialAttention::new(&vs / "height", planes, groups, false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn forward_t(&self, x: &Tensor, train: bool) -> Tensor {
|
||||
let x = self.width_axis.forward_t(x, train);
|
||||
self.height_axis.forward_t(&x, train)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
//! WiFlow-STD — spatio-temporal-decoupled CSI pose estimation (ADR-152 §2.2).
|
||||
//!
|
||||
//! Native Rust port of the **WiFlow-STD** architecture by DY2434
|
||||
//! (<https://github.com/DY2434/WiFlow-WiFi-Pose-Estimation-with-Spatio-Temporal-Decoupling>,
|
||||
//! Apache-2.0), reimplemented idiomatically from the vendored read-only
|
||||
//! reference in `benchmarks/wiflow-std/upstream/models/`.
|
||||
//!
|
||||
//! ## Evidence grade (ADR-152 §2.2 citation rule)
|
||||
//!
|
||||
//! Per `benchmarks/wiflow-std/RESULTS.md`, the upstream accuracy claims are
|
||||
//! **MEASURED-EQUIVALENT**: our retraining of the reference implementation on
|
||||
//! the released dataset reproduced **~96% PCK@20** (96.09% full test / 96.61%
|
||||
//! corruption-free; published claim 97.25%). The *shipped* upstream checkpoint
|
||||
//! was REFUTED (0.08% PCK@20 — keypoint-convention mismatch), and the released
|
||||
//! dataset/code required repairs before training converged. Cite this port as
|
||||
//! "~96% PCK@20 (our reproduction)" — **not comparable** to RuView's
|
||||
//! 17-keypoint ESP32 numbers (different hardware, subjects, split, skeleton).
|
||||
//!
|
||||
//! ## Name collision
|
||||
//!
|
||||
//! WiFlow-STD (this module) is the *external* DY2434 architecture. It is
|
||||
//! **distinct from RuView's internal WiFlow** camera-free pose pipeline; the
|
||||
//! `_std` suffix (Spatio-Temporal Decoupling) disambiguates the two.
|
||||
//!
|
||||
//! ## Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! CSI window [B, 540 sub, 20 t]
|
||||
//! │ TCN stack: 4 × grouped TemporalBlock (groups=20, k=3, dilation 1/2/4/8,
|
||||
//! │ depthwise-grouped + pointwise convs, causal Chomp1d padding)
|
||||
//! ▼ channels 540 → 540 → 440 → 340 → 240
|
||||
//! [B, 240, 20] ── transpose+unsqueeze ──► [B, 1, 20, 240] (image-like)
|
||||
//! │ ConvBlock1 (1→8, asymmetric 1×3 kernels, no downsampling)
|
||||
//! │ 4 × AsymmetricConvBlock (8→8→16→32→64, stride (1,2) on subcarrier axis)
|
||||
//! ▼
|
||||
//! [B, 64, 20, 15] ── permute ──► [B, 64, 15, 20]
|
||||
//! │ DualAxialAttention (64 ch, 8 groups, width- then height-axial
|
||||
//! │ self-attention with BN-normalised qkv and BN-normalised similarity)
|
||||
//! │ Decoder convs 64 → 32 → 2 (3×3 then 1×1, BN + SiLU)
|
||||
//! ▼
|
||||
//! [B, 2, 15, 20] ── adaptive avg-pool (K, 1) ──► [B, K, 2] keypoints
|
||||
//! ```
|
||||
//!
|
||||
//! 2,225,042 parameters / ~0.055 GFLOPs at the 15-keypoint default
|
||||
//! (both verified against the reference — see `RESULTS.md`).
|
||||
//!
|
||||
//! Note: upstream `config.py` lists `TCN_CHANNELS = [480, 360, 240]`, but the
|
||||
//! released checkpoint and `models/` code use `[540, 440, 340, 240]`. This
|
||||
//! port follows the `models/` code, which we verified loads the released
|
||||
//! weights after key remapping.
|
||||
//!
|
||||
//! ## Feature gating
|
||||
//!
|
||||
//! [`WiFlowStdConfig`] (validation, parameter-count formula, output-shape
|
||||
//! inference) is pure Rust and always available. [`model::WiFlowStdModel`]
|
||||
//! (the tch / LibTorch forward pass) requires the `tch-backend` feature,
|
||||
//! matching [`crate::model`]'s gating.
|
||||
|
||||
pub mod config;
|
||||
|
||||
#[cfg(feature = "tch-backend")]
|
||||
mod layers;
|
||||
#[cfg(feature = "tch-backend")]
|
||||
pub mod model;
|
||||
|
||||
pub use config::{TcnGroupsMode, WiFlowStdConfig};
|
||||
|
||||
#[cfg(feature = "tch-backend")]
|
||||
pub use model::WiFlowStdModel;
|
||||
@@ -0,0 +1,360 @@
|
||||
//! WiFlow-STD forward pass (tch-rs / LibTorch backend, ADR-152 §2.2).
|
||||
//!
|
||||
//! Idiomatic reimplementation of the DY2434 reference (Apache-2.0); see the
|
||||
//! [module docs](crate::wiflow_std) for provenance and the evidence grade.
|
||||
//! From-scratch init: BN gamma is pinned to 1 (see `layers::bn_cfg`); the
|
||||
//! axial-attention qkv conv uses `N(0, sqrt(1/in_planes))` per the
|
||||
//! reference's `attention.py` intent (note the reference's *effective* init
|
||||
//! differs — its `_initialize_weights` re-inits every `nn.Conv1d`, qkv
|
||||
//! included, with `kaiming_normal(fan_out)`); conv weights keep tch defaults
|
||||
//! (kaiming-uniform fan_in), which differ in scale from PyTorch's defaults.
|
||||
//! These divergences affect from-scratch training dynamics only — BN absorbs
|
||||
//! them at init, and loaded checkpoints overwrite everything. The
|
||||
//! retrained PyTorch checkpoint loads via [`WiFlowStdModel::load`] after
|
||||
//! key-remapped safetensors export
|
||||
//! (`benchmarks/wiflow-std/export_to_safetensors.py`); numerical parity with
|
||||
//! the PyTorch forward pass is proven by
|
||||
//! `tests/test_wiflow_std_parity.rs` (max abs diff ~1.2e-7).
|
||||
|
||||
use tch::{nn, Device, Tensor};
|
||||
|
||||
use super::config::WiFlowStdConfig;
|
||||
use super::layers::{ConvBlock, DualAxialAttention, GroupedTemporalBlock};
|
||||
use crate::error::TrainError;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// WiFlowStdModel
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// WiFlow-STD pose model: TCN temporal encoder → asymmetric 2-D conv encoder
|
||||
/// → dual axial attention → conv decoder → adaptive pool to `(K, 2)` keypoints.
|
||||
///
|
||||
/// Input: `[B, subcarriers, window]` CSI amplitudes.
|
||||
/// Output: `[B, keypoints, 2]` normalised 2-D keypoint coordinates.
|
||||
pub struct WiFlowStdModel {
|
||||
vs: nn::VarStore,
|
||||
tcn: Vec<GroupedTemporalBlock>,
|
||||
conv_in: ConvBlock,
|
||||
conv_blocks: Vec<ConvBlock>,
|
||||
attention: DualAxialAttention,
|
||||
dec_conv1: nn::Conv2D,
|
||||
dec_bn1: nn::BatchNorm,
|
||||
dec_conv2: nn::Conv2D,
|
||||
dec_bn2: nn::BatchNorm,
|
||||
/// Active model configuration.
|
||||
pub config: WiFlowStdConfig,
|
||||
}
|
||||
|
||||
impl WiFlowStdModel {
|
||||
/// Build a new model with randomly-initialised weights on `device`.
|
||||
///
|
||||
/// Call `tch::manual_seed(seed)` before this for reproducibility.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`TrainError::Config`] if `config.validate()` fails.
|
||||
pub fn new(config: &WiFlowStdConfig, device: Device) -> Result<Self, TrainError> {
|
||||
config.validate()?;
|
||||
|
||||
let vs = nn::VarStore::new(device);
|
||||
let root = vs.root();
|
||||
|
||||
// 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;
|
||||
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 as i64,
|
||||
c_out as i64,
|
||||
dilation,
|
||||
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;
|
||||
}
|
||||
|
||||
// 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,
|
||||
strides[i] as i64,
|
||||
));
|
||||
c_in = c_out as i64;
|
||||
}
|
||||
|
||||
let attention =
|
||||
DualAxialAttention::new(&root / "attention", c_in, config.attention_groups as i64);
|
||||
|
||||
// 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,
|
||||
mid,
|
||||
3,
|
||||
nn::ConvConfig {
|
||||
padding: 1,
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let dec_bn1 = nn::batch_norm2d(&root / "dec_bn1", mid, super::layers::bn_cfg());
|
||||
let dec_conv2 = nn::conv2d(&root / "dec_conv2", mid, 2, 1, Default::default());
|
||||
let dec_bn2 = nn::batch_norm2d(&root / "dec_bn2", 2, super::layers::bn_cfg());
|
||||
|
||||
Ok(WiFlowStdModel {
|
||||
vs,
|
||||
tcn,
|
||||
conv_in,
|
||||
conv_blocks,
|
||||
attention,
|
||||
dec_conv1,
|
||||
dec_bn1,
|
||||
dec_conv2,
|
||||
dec_bn2,
|
||||
config: config.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass in training mode (dropout active, BN in train mode).
|
||||
///
|
||||
/// `csi`: `[B, subcarriers, window]` → `[B, keypoints, 2]`.
|
||||
pub fn forward_t(&self, csi: &Tensor) -> Tensor {
|
||||
self.forward_impl(csi, true)
|
||||
}
|
||||
|
||||
/// Forward pass without gradient tracking (inference mode).
|
||||
pub fn forward_inference(&self, csi: &Tensor) -> Tensor {
|
||||
tch::no_grad(|| self.forward_impl(csi, false))
|
||||
}
|
||||
|
||||
/// Save model weights. The tch `VarStore` dispatches the format on the
|
||||
/// file extension: `.safetensors` → safetensors, anything else → torch
|
||||
/// `.pt`.
|
||||
///
|
||||
/// **Platform constraint:** prefer `.safetensors`. The `.pt` path
|
||||
/// (`_save_parameters`/`_load_parameters`) is broken on Windows with
|
||||
/// torch 2.11 (GenericDict internal assert on the load roundtrip — see
|
||||
/// the `save_and_load_roundtrip` test below), and the verified retrained
|
||||
/// checkpoint is shipped as key-remapped safetensors anyway
|
||||
/// (`benchmarks/wiflow-std/export_to_safetensors.py`).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`TrainError::TrainingStep`] if the file cannot be written.
|
||||
pub fn save(&self, path: &std::path::Path) -> Result<(), TrainError> {
|
||||
self.vs
|
||||
.save(path)
|
||||
.map_err(|e| TrainError::training_step(format!("save failed: {e}")))
|
||||
}
|
||||
|
||||
/// Load model weights from a file (format dispatched on extension; see
|
||||
/// the `.pt`-on-Windows caveat on [`Self::save`]).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`TrainError::TrainingStep`] if the file cannot be read or the
|
||||
/// weights are incompatible with this architecture.
|
||||
pub fn load(&mut self, path: &std::path::Path) -> Result<(), TrainError> {
|
||||
self.vs
|
||||
.load(path)
|
||||
.map_err(|e| TrainError::training_step(format!("load failed: {e}")))
|
||||
}
|
||||
|
||||
/// Reference to the internal `VarStore` (e.g. to build an optimiser).
|
||||
pub fn var_store(&self) -> &nn::VarStore {
|
||||
&self.vs
|
||||
}
|
||||
|
||||
/// Mutable access to the internal `VarStore`.
|
||||
pub fn var_store_mut(&mut self) -> &mut nn::VarStore {
|
||||
&mut self.vs
|
||||
}
|
||||
|
||||
/// Total number of trainable scalar parameters. Must equal
|
||||
/// [`WiFlowStdConfig::param_count`] (2,225,042 at the default config).
|
||||
pub fn num_parameters(&self) -> i64 {
|
||||
self.vs
|
||||
.trainable_variables()
|
||||
.iter()
|
||||
.map(|t| t.numel() as i64)
|
||||
.sum()
|
||||
}
|
||||
|
||||
fn forward_impl(&self, csi: &Tensor, train: bool) -> Tensor {
|
||||
// TCN: [B, subcarriers, T] → [B, c_tcn, T].
|
||||
let mut h = csi.shallow_clone();
|
||||
for block in &self.tcn {
|
||||
h = block.forward_t(&h, train);
|
||||
}
|
||||
|
||||
// Image-like reshape: [B, c_tcn, T] → [B, 1, T, c_tcn].
|
||||
let h = h.transpose(1, 2).unsqueeze(1);
|
||||
|
||||
// 2-D conv encoder: [B, 1, T, S] → [B, C, T, S'].
|
||||
let mut h = self.conv_in.forward_t(&h, train);
|
||||
for block in &self.conv_blocks {
|
||||
h = block.forward_t(&h, train);
|
||||
}
|
||||
|
||||
// Swap to [B, C, S', T] for the axial attention + decoder.
|
||||
let h = h.permute([0, 1, 3, 2]);
|
||||
let h = self.attention.forward_t(&h, train);
|
||||
|
||||
// Decoder: [B, C, S', T] → [B, 2, S', T].
|
||||
let h = h
|
||||
.apply(&self.dec_conv1)
|
||||
.apply_t(&self.dec_bn1, train)
|
||||
.silu()
|
||||
.apply(&self.dec_conv2)
|
||||
.apply_t(&self.dec_bn2, train)
|
||||
.silu();
|
||||
|
||||
// [B, 2, S', T] → pool (K, 1) → [B, 2, K] → [B, K, 2].
|
||||
let k = self.config.keypoints as i64;
|
||||
h.adaptive_avg_pool2d([k, 1])
|
||||
.squeeze_dim(-1)
|
||||
.transpose(1, 2)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests (require the tch-backend feature + LibTorch)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tch::Kind;
|
||||
|
||||
fn random_csi(cfg: &WiFlowStdConfig, batch: i64) -> Tensor {
|
||||
Tensor::rand(
|
||||
[batch, cfg.subcarriers as i64, cfg.window as i64],
|
||||
(Kind::Float, Device::Cpu),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn param_count_matches_pure_rust_formula() {
|
||||
tch::manual_seed(0);
|
||||
let cfg = WiFlowStdConfig::default();
|
||||
let model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("default config builds");
|
||||
// Pins the tch graph against the verified reference (2,225,042).
|
||||
assert_eq!(model.num_parameters(), cfg.param_count() as i64);
|
||||
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);
|
||||
let cfg = WiFlowStdConfig::default();
|
||||
let model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("build");
|
||||
let out = model.forward_t(&random_csi(&cfg, 2));
|
||||
assert_eq!(out.size(), &[2, 15, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_output_shape_17_keypoints_esp32() {
|
||||
tch::manual_seed(0);
|
||||
let cfg = WiFlowStdConfig::for_keypoints(17);
|
||||
let model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("build");
|
||||
let out = model.forward_inference(&random_csi(&cfg, 1));
|
||||
assert_eq!(out.size(), &[1, 17, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inference_outputs_are_finite_and_deterministic() {
|
||||
tch::manual_seed(7);
|
||||
let cfg = WiFlowStdConfig::default();
|
||||
let model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("build");
|
||||
let csi = random_csi(&cfg, 1);
|
||||
let a = model.forward_inference(&csi);
|
||||
let b = model.forward_inference(&csi);
|
||||
assert!(
|
||||
bool::try_from(a.isfinite().all()).unwrap(),
|
||||
"non-finite output"
|
||||
);
|
||||
assert!(
|
||||
bool::try_from(a.eq_tensor(&b).all()).unwrap(),
|
||||
"inference must be deterministic (dropout disabled)"
|
||||
);
|
||||
}
|
||||
|
||||
/// Dumps the authoritative tch `VarStore` variable names + shapes. This is
|
||||
/// the source of truth for the PyTorch→tch key mapping implemented in
|
||||
/// `benchmarks/wiflow-std/export_to_safetensors.py` — rerun it (with
|
||||
/// `--nocapture`) whenever the architecture changes.
|
||||
#[test]
|
||||
fn dump_variable_names() {
|
||||
let cfg = WiFlowStdConfig::default();
|
||||
let model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("build");
|
||||
let vars = model.var_store().variables();
|
||||
let mut names: Vec<(String, Vec<i64>)> =
|
||||
vars.iter().map(|(n, t)| (n.clone(), t.size())).collect();
|
||||
names.sort();
|
||||
for (name, shape) in &names {
|
||||
println!("{name} {shape:?}");
|
||||
}
|
||||
println!("total: {} variables", names.len());
|
||||
assert!(!names.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_config_is_rejected() {
|
||||
let cfg = WiFlowStdConfig {
|
||||
subcarriers: 541, // not divisible by tcn_groups
|
||||
..Default::default()
|
||||
};
|
||||
assert!(WiFlowStdModel::new(&cfg, Device::Cpu).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn save_and_load_roundtrip() {
|
||||
use tempfile::tempdir;
|
||||
tch::manual_seed(42);
|
||||
let cfg = WiFlowStdConfig::default();
|
||||
let mut model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("build");
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
// safetensors, not .pt: this torch build's _save_parameters/_load_parameters
|
||||
// .pt roundtrip is broken on Windows (GenericDict internal assert)
|
||||
let path = tmp.path().join("wiflow_std.safetensors");
|
||||
model.save(&path).expect("save");
|
||||
model.load(&path).expect("load");
|
||||
let out = model.forward_inference(&random_csi(&cfg, 1));
|
||||
assert_eq!(out.size(), &[1, 15, 2]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user