frank/ruvnet--RuView
1
0
Fork 0
mirror of https://github.com/ruvnet/RuView synced 2026-08-10 20:31:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: ADR-152 Rust integrations + ADR-153 802.11bf protocol model

Browse Source 
- 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>
This commit is contained in:
ruv
2026-06-10 19:46:34 -04:00
parent e1936c9a24
commit b0c324edc5
31 changed files with 5024 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
v2/crates/wifi-densepose-train/src/error.rs
+76 -1
View File
@@ -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,73 @@ 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,
},
/// 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,
},
}
v2/crates/wifi-densepose-train/src/lib.rs
+11 -1
View File
@@ -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.
@@ -78,7 +80,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;
@@ -86,6 +88,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;
v2/crates/wifi-densepose-train/src/mae.rs
+379
View File
@@ -0,0 +1,379 @@
//! 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`].
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`.
#[must_use]
pub fn random_mask(n_patches: usize, mask_ratio: f64, seed: u64) -> MaskIndices {
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();
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(())
}
v2/crates/wifi-densepose-train/src/wiflow_std/config.rs
+476
View File
@@ -0,0 +1,476 @@
//! 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;
// ---------------------------------------------------------------------------
// 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**.
pub tcn_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.
pub keypoints: 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,
conv_channels: vec![8, 16, 32, 64],
attention_groups: 8,
keypoints: 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)).
pub fn for_keypoints(keypoints: usize) -> Self {
WiFlowStdConfig {
keypoints,
..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"));
}
if 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 || 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.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.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)
}
/// Width of the encoder feature map after the strided conv blocks.
///
/// `ConvBlock1` preserves width; each `AsymmetricConvBlock` applies a
/// `(1, 3)` kernel with stride `(1, 2)` and padding `(0, 1)`:
/// `w → (w - 1) / 2 + 1`. Default: 240 → 120 → 60 → 30 → **15**.
pub fn feature_width(&self) -> usize {
let mut w = self.tcn_output_channels();
for _ in &self.conv_channels {
w = (w.saturating_sub(1)) / 2 + 1;
}
w
}
/// 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).
pub fn param_count(&self) -> usize {
let mut total = 0;
// TCN stack.
let mut c_in = self.subcarriers;
for &c_out in &self.tcn_channels {
total += tcn_block_params(c_in, c_out, TCN_KERNEL, self.tcn_groups);
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 → c/2) + BN + 1×1 conv (c/2 → 2) + BN.
total += decoder_params(c_in);
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.
fn tcn_block_params(c_in: usize, c_out: usize, k: usize, groups: usize) -> usize {
let grouped1 = c_in * (c_in / groups) * k; // depthwise-grouped, c_in → c_in
let bn1g = 2 * c_in;
let pw1 = c_out * c_in; // pointwise 1×1
let bn1p = 2 * c_out;
let grouped2 = c_out * (c_out / groups) * k;
let bn2g = 2 * c_out;
let pw2 = c_out * c_out;
let bn2p = 2 * c_out;
let downsample = if c_in != c_out {
c_in * c_out + 2 * c_out
} 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 → c/2, 3×3, bias)` + BN + `Conv2d(c/2 → 2, 1×1, bias)`
/// + BN.
fn decoder_params(c: usize) -> usize {
let mid = c / 2;
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), 675_000);
assert_eq!(tcn_block_params(540, 440, 3, 20), 746_180);
assert_eq!(tcn_block_params(440, 340, 3, 20), 464_780);
assert_eq!(tcn_block_params(340, 240, 3, 20), 249_380);
// 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), 18_598);
}
#[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);
}
}
v2/crates/wifi-densepose-train/src/wiflow_std/layers.rs
+314
View File
@@ -0,0 +1,314 @@
//! 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};
// ---------------------------------------------------------------------------
// 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 {
pub(super) fn new(
vs: nn::Path,
c_in: i64,
c_out: i64,
dilation: i64,
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 = nn::ConvConfig {
bias: false,
..Default::default()
};
let conv1_group = nn::conv1d(&vs / "conv1_group", c_in, c_in, k, grouped_cfg(groups));
let bn1_group = nn::batch_norm1d(&vs / "bn1_group", c_in, Default::default());
let conv1_pw = nn::conv1d(&vs / "conv1_pw", c_in, c_out, 1, pointwise_cfg);
let bn1_pw = nn::batch_norm1d(&vs / "bn1_pw", c_out, Default::default());
let conv2_group = nn::conv1d(&vs / "conv2_group", c_out, c_out, k, grouped_cfg(groups));
let bn2_group = nn::batch_norm1d(&vs / "bn2_group", c_out, Default::default());
let conv2_pw = nn::conv1d(&vs / "conv2_pw", c_out, c_out, 1, pointwise_cfg);
let bn2_pw = nn::batch_norm1d(&vs / "bn2_pw", c_out, Default::default());
let downsample = (c_in != c_out).then(|| {
(
nn::conv1d(&vs / "ds_conv", c_in, c_out, 1, pointwise_cfg),
nn::batch_norm1d(&vs / "ds_bn", c_out, Default::default()),
)
});
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, Default::default());
let conv2 = nn::conv(&vs / "conv2", c_out, c_out, [1, 3], asym(1));
let bn2 = nn::batch_norm2d(&vs / "bn2", c_out, Default::default());
let conv3 = nn::conv(&vs / "conv3", c_out, c_out, [1, 3], asym(1));
let bn3 = nn::batch_norm2d(&vs / "bn3", c_out, Default::default());
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, Default::default());
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, Default::default());
let bn_similarity = nn::batch_norm2d(&vs / "bn_similarity", groups, Default::default());
let bn_output = nn::batch_norm1d(&vs / "bn_output", planes, Default::default());
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)
}
}
v2/crates/wifi-densepose-train/src/wiflow_std/mod.rs
+69
View File
@@ -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::WiFlowStdConfig;
#[cfg(feature = "tch-backend")]
pub use model::WiFlowStdModel;
v2/crates/wifi-densepose-train/src/wiflow_std/model.rs
+292
View File
@@ -0,0 +1,292 @@
//! 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.
//! Weights are initialised from scratch (tch defaults; the axial-attention
//! qkv conv mirrors the reference's `N(0, sqrt(1/in_planes))` init). Loading
//! the retrained PyTorch checkpoint is a follow-up (key remap + `vs.load`).
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.
let mut tcn = Vec::with_capacity(config.tcn_channels.len());
let mut c_in = config.subcarriers as i64;
for (i, &c_out) in config.tcn_channels.iter().enumerate() {
let dilation = 1_i64 << i;
tcn.push(GroupedTemporalBlock::new(
&root / format!("tcn{i}"),
c_in,
c_out as i64,
dilation,
config.tcn_groups as i64,
config.dropout,
));
c_in = c_out as i64;
}
// 2-D conv encoder: ConvBlock1 (stride 1) + strided asymmetric blocks.
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 mut c_in = c0;
for (i, &c_out) in config.conv_channels.iter().enumerate() {
conv_blocks.push(ConvBlock::new(
&root / format!("conv{i}"),
c_in,
c_out as i64,
2,
));
c_in = c_out as i64;
}
let attention =
DualAxialAttention::new(&root / "attention", c_in, config.attention_groups as i64);
// Decoder: c → c/2 (3×3) → 2 (1×1), BN + SiLU after each conv.
let mid = c_in / 2;
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, Default::default());
let dec_conv2 = nn::conv2d(&root / "dec_conv2", mid, 2, 1, Default::default());
let dec_bn2 = nn::batch_norm2d(&root / "dec_bn2", 2, Default::default());
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 (tch `.pt` / safetensors format).
///
/// # 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.
///
/// # 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);
}
#[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)"
);
}
#[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");
let path = tmp.path().join("wiflow_std.pt");
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]);
}
}