mirror of
https://github.com/ruvnet/RuView
synced 2026-07-25 17:51:48 +00:00
feat(train): compact WiFlow-STD presets in Rust + tiny edge artifact (ADR-152)
WiFlowStdConfig gains half()/quarter()/tiny() mirroring the overnight sweep exactly: TcnGroupsMode (Fixed/Gcd/Depthwise), input_pw_groups, derived stride schedule and decoder-mid (all default to upstream behavior; legacy serde JSON unaffected). Param formulas pin to trained ground truth first try: 843,834 / 338,600 / 56,290; default 2,225,042 pin and 1.192e-7 parity unchanged. 248 tests green. Tiny edge artifact (tiny_edge_bench.py): ONNX fp32 = 295 KB, 0.66 ms/win (~1,500/s CPU), 94.11% PCK@20 (matches sweep clean-test exactly; parity 1.49e-7). Static int8 is a bad trade at this scale (-1.43pt, +19% MPJPE, -16% size, slower) — recorded as negative result. Export note: width-16 breaks AdaptiveAvgPool((15,1)) TorchScript export; replaced by exact mean+matmul equivalent, proven by parity. Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruv
@@ -19,6 +19,43 @@ pub const TCN_KERNEL: usize = 3;
|
||||
/// forwarded to the TCN), so it is a constant here rather than a config field.
|
||||
pub const CONV_BLOCK_DROPOUT: f64 = 0.3;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TcnGroupsMode
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// How the group count of each depthwise-grouped TCN convolution is chosen
|
||||
/// (ADR-152 efficiency sweep, `benchmarks/wiflow-std/remote/sweep/model_compact.py`).
|
||||
///
|
||||
/// The upstream reference hardcodes `groups = 20`, which does not divide the
|
||||
/// compact variants' channel counts (e.g. 270, 135, 85). The sweep's rules:
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TcnGroupsMode {
|
||||
/// Every grouped conv uses [`WiFlowStdConfig::tcn_groups`] verbatim
|
||||
/// (upstream behavior; requires divisibility). Default.
|
||||
#[default]
|
||||
Fixed,
|
||||
/// Per-conv groups = `gcd(channels, tcn_groups)` — equals `tcn_groups`
|
||||
/// wherever the upstream choice is valid (incl. the 540-channel input
|
||||
/// conv) and falls back to the largest common divisor otherwise.
|
||||
/// The sweep's `gcd20` mode (`half` / `quarter` presets).
|
||||
Gcd,
|
||||
/// Per-conv groups = channels (fully depthwise; `tiny` preset).
|
||||
Depthwise,
|
||||
}
|
||||
|
||||
fn gcd(a: usize, b: usize) -> usize {
|
||||
let (mut a, mut b) = (a, b);
|
||||
while b != 0 {
|
||||
(a, b) = (b, a % b);
|
||||
}
|
||||
a
|
||||
}
|
||||
|
||||
fn default_input_pw_groups() -> usize {
|
||||
1
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// WiFlowStdConfig
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -47,9 +84,26 @@ pub struct WiFlowStdConfig {
|
||||
|
||||
/// Group count for the depthwise-grouped TCN convolutions. The reference
|
||||
/// hardcodes **20**; exposed so non-540 subcarrier layouts can keep the
|
||||
/// divisibility invariant. Default: **20**.
|
||||
/// divisibility invariant. Default: **20**. Interpreted per
|
||||
/// [`Self::tcn_groups_mode`]: the verbatim group count in `Fixed` mode,
|
||||
/// the gcd base in `Gcd` mode, ignored in `Depthwise` mode.
|
||||
pub tcn_groups: usize,
|
||||
|
||||
/// Group-selection rule for the TCN's grouped convolutions
|
||||
/// (ADR-152 efficiency sweep). Default: [`TcnGroupsMode::Fixed`]
|
||||
/// (upstream behavior — every grouped conv uses [`Self::tcn_groups`]).
|
||||
#[serde(default)]
|
||||
pub tcn_groups_mode: TcnGroupsMode,
|
||||
|
||||
/// Group count for the **first** TCN block's pointwise (1×1) and residual
|
||||
/// downsample convs (`subcarriers → tcn_channels[0]`). The sweep's `tiny`
|
||||
/// variant uses **4** to break the dense-540-input parameter floor
|
||||
/// (~117k params, which alone exceeds tiny's budget); every other config
|
||||
/// uses **1** (upstream behavior). Must divide both `subcarriers` and
|
||||
/// `tcn_channels[0]`. Default: **1**.
|
||||
#[serde(default = "default_input_pw_groups")]
|
||||
pub input_pw_groups: usize,
|
||||
|
||||
/// Output channels of the 2-D conv encoder blocks. The first entry is
|
||||
/// also `ConvBlock1`'s output; each subsequent block downsamples the
|
||||
/// subcarrier axis by 2. Default: **[8, 16, 32, 64]**.
|
||||
@@ -75,6 +129,8 @@ impl Default for WiFlowStdConfig {
|
||||
window: 20,
|
||||
tcn_channels: vec![540, 440, 340, 240],
|
||||
tcn_groups: 20,
|
||||
tcn_groups_mode: TcnGroupsMode::Fixed,
|
||||
input_pw_groups: 1,
|
||||
conv_channels: vec![8, 16, 32, 64],
|
||||
attention_groups: 8,
|
||||
keypoints: 15,
|
||||
@@ -93,6 +149,52 @@ impl WiFlowStdConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// **half** compact preset (ADR-152 efficiency sweep, trained
|
||||
/// 2026-06-10/11): **843,834** parameters (0.38×), clean-test PCK@20
|
||||
/// **96.62%** — strictly dominates the full reference on its own
|
||||
/// benchmark. Per-conv groups = `gcd(channels, 20)`; stride schedule
|
||||
/// derives to `[2, 2, 2, 1]`. See
|
||||
/// `benchmarks/wiflow-std/results/efficiency_sweep.jsonl`.
|
||||
pub fn half() -> Self {
|
||||
WiFlowStdConfig {
|
||||
tcn_channels: vec![270, 220, 170, 120],
|
||||
tcn_groups_mode: TcnGroupsMode::Gcd,
|
||||
conv_channels: vec![4, 8, 16, 32],
|
||||
attention_groups: 4,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// **quarter** compact preset (ADR-152 efficiency sweep): **338,600**
|
||||
/// parameters (0.15×), clean-test PCK@20 **96.05%**. Per-conv groups =
|
||||
/// `gcd(channels, 20)`; stride schedule derives to `[2, 2, 1, 1]`.
|
||||
pub fn quarter() -> Self {
|
||||
WiFlowStdConfig {
|
||||
tcn_channels: vec![135, 110, 85, 60],
|
||||
tcn_groups_mode: TcnGroupsMode::Gcd,
|
||||
conv_channels: vec![2, 4, 8, 16],
|
||||
attention_groups: 2,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// **tiny** compact preset (ADR-152 efficiency sweep): **56,290**
|
||||
/// parameters (0.025×), clean-test PCK@20 **94.11%** — the smallest
|
||||
/// deployable WiFlow-class model (~220 KB fp32). Fully depthwise TCN
|
||||
/// groups plus `input_pw_groups = 4` on the first block's pointwise /
|
||||
/// downsample convs; stride schedule derives to `[2, 1, 1, 1]`
|
||||
/// (feature width 16).
|
||||
pub fn tiny() -> Self {
|
||||
WiFlowStdConfig {
|
||||
tcn_channels: vec![68, 56, 44, 32],
|
||||
tcn_groups_mode: TcnGroupsMode::Depthwise,
|
||||
input_pw_groups: 4,
|
||||
conv_channels: vec![2, 4, 8, 16],
|
||||
attention_groups: 2,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate all architectural invariants.
|
||||
///
|
||||
/// # Errors
|
||||
@@ -108,7 +210,12 @@ impl WiFlowStdConfig {
|
||||
if self.tcn_groups == 0 {
|
||||
return Err(ConfigError::invalid_value("tcn_groups", "must be >= 1"));
|
||||
}
|
||||
if self.subcarriers % self.tcn_groups != 0 {
|
||||
// In Gcd mode the per-conv group count is gcd(channels, tcn_groups)
|
||||
// and in Depthwise mode it is the channel count itself, so the
|
||||
// divisibility invariant holds by construction; only Fixed mode
|
||||
// (upstream behavior) needs the explicit checks.
|
||||
let fixed = self.tcn_groups_mode == TcnGroupsMode::Fixed;
|
||||
if fixed && self.subcarriers % self.tcn_groups != 0 {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"subcarriers",
|
||||
format!(
|
||||
@@ -124,7 +231,7 @@ impl WiFlowStdConfig {
|
||||
));
|
||||
}
|
||||
for (i, &c) in self.tcn_channels.iter().enumerate() {
|
||||
if c == 0 || c % self.tcn_groups != 0 {
|
||||
if c == 0 || (fixed && c % self.tcn_groups != 0) {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"tcn_channels",
|
||||
format!(
|
||||
@@ -134,6 +241,18 @@ impl WiFlowStdConfig {
|
||||
));
|
||||
}
|
||||
}
|
||||
if self.input_pw_groups == 0
|
||||
|| self.subcarriers % self.input_pw_groups != 0
|
||||
|| self.tcn_channels[0] % self.input_pw_groups != 0
|
||||
{
|
||||
return Err(ConfigError::invalid_value(
|
||||
"input_pw_groups",
|
||||
format!(
|
||||
"{} must be >= 1 and divide both subcarriers={} and tcn_channels[0]={}",
|
||||
self.input_pw_groups, self.subcarriers, self.tcn_channels[0]
|
||||
),
|
||||
));
|
||||
}
|
||||
if self.conv_channels.is_empty() {
|
||||
return Err(ConfigError::invalid_value(
|
||||
"conv_channels",
|
||||
@@ -184,19 +303,61 @@ impl WiFlowStdConfig {
|
||||
*self.tcn_channels.last().unwrap_or(&0)
|
||||
}
|
||||
|
||||
/// Width of the encoder feature map after the strided conv blocks.
|
||||
/// Group count of a grouped TCN conv over `channels` channels, per
|
||||
/// [`Self::tcn_groups_mode`].
|
||||
pub fn tcn_conv_groups(&self, channels: usize) -> usize {
|
||||
match self.tcn_groups_mode {
|
||||
TcnGroupsMode::Fixed => self.tcn_groups,
|
||||
TcnGroupsMode::Gcd => gcd(channels, self.tcn_groups),
|
||||
TcnGroupsMode::Depthwise => channels,
|
||||
}
|
||||
}
|
||||
|
||||
/// Width stride of each `AsymmetricConvBlock`, derived with the sweep's
|
||||
/// rule (`model_compact.py::compute_strides`): halve the width
|
||||
/// (`w → ceil(w / 2)`, the `(1,3)`-kernel stride-2 output size) only
|
||||
/// while the result stays ≥ [`Self::keypoints`], so the final adaptive
|
||||
/// pool never has to duplicate rows. At the upstream default
|
||||
/// (240 channels, 15 keypoints) this derives `[2, 2, 2, 2]` — the
|
||||
/// hardcoded upstream schedule, exactly.
|
||||
pub fn conv_strides(&self) -> Vec<usize> {
|
||||
let mut w = self.tcn_output_channels();
|
||||
let mut strides = Vec::with_capacity(self.conv_channels.len());
|
||||
for _ in &self.conv_channels {
|
||||
let next = w.div_ceil(2);
|
||||
if next >= self.keypoints {
|
||||
strides.push(2);
|
||||
w = next;
|
||||
} else {
|
||||
strides.push(1);
|
||||
}
|
||||
}
|
||||
strides
|
||||
}
|
||||
|
||||
/// Width of the encoder feature map after the conv blocks.
|
||||
///
|
||||
/// `ConvBlock1` preserves width; each `AsymmetricConvBlock` applies a
|
||||
/// `(1, 3)` kernel with stride `(1, 2)` and padding `(0, 1)`:
|
||||
/// `w → (w - 1) / 2 + 1`. Default: 240 → 120 → 60 → 30 → **15**.
|
||||
/// `(1, 3)` kernel with padding `(0, 1)` and the per-block stride from
|
||||
/// [`Self::conv_strides`]. Default: 240 → 120 → 60 → 30 → **15**.
|
||||
pub fn feature_width(&self) -> usize {
|
||||
let mut w = self.tcn_output_channels();
|
||||
for _ in &self.conv_channels {
|
||||
w = (w.saturating_sub(1)) / 2 + 1;
|
||||
for s in self.conv_strides() {
|
||||
if s == 2 {
|
||||
w = w.div_ceil(2);
|
||||
}
|
||||
}
|
||||
w
|
||||
}
|
||||
|
||||
/// Mid-channel count of the decoder's 3×3 conv:
|
||||
/// `max(conv_channels.last() / 2, 4)` (the sweep's floor of 4 keeps the
|
||||
/// decoder viable at very small widths; identical to the upstream `c / 2`
|
||||
/// for every channel count ≥ 8, including the default 64 → 32).
|
||||
pub fn decoder_mid(&self) -> usize {
|
||||
(self.conv_channels.last().unwrap_or(&0) / 2).max(4)
|
||||
}
|
||||
|
||||
/// Output tensor shape `(batch, keypoints, 2)`. The adaptive average pool
|
||||
/// maps the feature height to `keypoints` regardless of its size, so the
|
||||
/// keypoint count is free (15 and 17 share identical weights).
|
||||
@@ -217,10 +378,19 @@ impl WiFlowStdConfig {
|
||||
pub fn param_count(&self) -> usize {
|
||||
let mut total = 0;
|
||||
|
||||
// TCN stack.
|
||||
// TCN stack: per-conv groups follow tcn_groups_mode; only the first
|
||||
// block's pointwise/downsample convs use input_pw_groups.
|
||||
let mut c_in = self.subcarriers;
|
||||
for &c_out in &self.tcn_channels {
|
||||
total += tcn_block_params(c_in, c_out, TCN_KERNEL, self.tcn_groups);
|
||||
for (i, &c_out) in self.tcn_channels.iter().enumerate() {
|
||||
let pw_groups = if i == 0 { self.input_pw_groups } else { 1 };
|
||||
total += tcn_block_params(
|
||||
c_in,
|
||||
c_out,
|
||||
TCN_KERNEL,
|
||||
self.tcn_conv_groups(c_in),
|
||||
self.tcn_conv_groups(c_out),
|
||||
pw_groups,
|
||||
);
|
||||
c_in = c_out;
|
||||
}
|
||||
|
||||
@@ -237,8 +407,8 @@ impl WiFlowStdConfig {
|
||||
// Dual axial attention: width axis + height axis, both c_in → c_in.
|
||||
total += 2 * axial_attention_params(c_in, self.attention_groups);
|
||||
|
||||
// Decoder: 3×3 conv (c → c/2) + BN + 1×1 conv (c/2 → 2) + BN.
|
||||
total += decoder_params(c_in);
|
||||
// Decoder: 3×3 conv (c → decoder_mid) + BN + 1×1 conv (mid → 2) + BN.
|
||||
total += decoder_params(c_in, self.decoder_mid());
|
||||
|
||||
total
|
||||
}
|
||||
@@ -250,18 +420,29 @@ impl WiFlowStdConfig {
|
||||
|
||||
/// One `InnerGroupedTemporalBlock`: two (depthwise-grouped conv → BN →
|
||||
/// pointwise conv → BN) stages plus a 1×1 + BN residual projection when the
|
||||
/// channel count changes. All convs are bias-free.
|
||||
fn tcn_block_params(c_in: usize, c_out: usize, k: usize, groups: usize) -> usize {
|
||||
let grouped1 = c_in * (c_in / groups) * k; // depthwise-grouped, c_in → c_in
|
||||
/// channel count changes. All convs are bias-free. `g_in`/`g_out` are the
|
||||
/// group counts of the two grouped convs (each conv groups over its own
|
||||
/// channel count — they differ in `Gcd`/`Depthwise` mode); `pw_groups`
|
||||
/// groups the first pointwise conv and the residual projection (the sweep's
|
||||
/// `input_pw_groups`, block 0 only — 1 everywhere else).
|
||||
fn tcn_block_params(
|
||||
c_in: usize,
|
||||
c_out: usize,
|
||||
k: usize,
|
||||
g_in: usize,
|
||||
g_out: usize,
|
||||
pw_groups: usize,
|
||||
) -> usize {
|
||||
let grouped1 = c_in * (c_in / g_in) * k; // depthwise-grouped, c_in → c_in
|
||||
let bn1g = 2 * c_in;
|
||||
let pw1 = c_out * c_in; // pointwise 1×1
|
||||
let pw1 = c_out * (c_in / pw_groups); // pointwise 1×1
|
||||
let bn1p = 2 * c_out;
|
||||
let grouped2 = c_out * (c_out / groups) * k;
|
||||
let grouped2 = c_out * (c_out / g_out) * k;
|
||||
let bn2g = 2 * c_out;
|
||||
let pw2 = c_out * c_out;
|
||||
let bn2p = 2 * c_out;
|
||||
let downsample = if c_in != c_out {
|
||||
c_in * c_out + 2 * c_out
|
||||
(c_in / pw_groups) * c_out + 2 * c_out
|
||||
} else {
|
||||
0
|
||||
};
|
||||
@@ -288,10 +469,9 @@ fn axial_attention_params(c: usize, groups: usize) -> usize {
|
||||
qkv + bn_qkv + bn_similarity + bn_output
|
||||
}
|
||||
|
||||
/// Decoder: `Conv2d(c → c/2, 3×3, bias)` + BN + `Conv2d(c/2 → 2, 1×1, bias)`
|
||||
/// + BN.
|
||||
fn decoder_params(c: usize) -> usize {
|
||||
let mid = c / 2;
|
||||
/// Decoder: `Conv2d(c → mid, 3×3, bias)` + BN + `Conv2d(mid → 2, 1×1, bias)`
|
||||
/// + BN, where `mid` = [`WiFlowStdConfig::decoder_mid`].
|
||||
fn decoder_params(c: usize, mid: usize) -> usize {
|
||||
let conv1 = mid * c * 9 + mid;
|
||||
let bn1 = 2 * mid;
|
||||
let conv2 = 2 * mid + 2;
|
||||
@@ -335,10 +515,10 @@ mod tests {
|
||||
#[test]
|
||||
fn per_component_breakdown_matches_hand_calculation() {
|
||||
// TCN levels (hand-verified against the reference layer shapes).
|
||||
assert_eq!(tcn_block_params(540, 540, 3, 20), 675_000);
|
||||
assert_eq!(tcn_block_params(540, 440, 3, 20), 746_180);
|
||||
assert_eq!(tcn_block_params(440, 340, 3, 20), 464_780);
|
||||
assert_eq!(tcn_block_params(340, 240, 3, 20), 249_380);
|
||||
assert_eq!(tcn_block_params(540, 540, 3, 20, 20, 1), 675_000);
|
||||
assert_eq!(tcn_block_params(540, 440, 3, 20, 20, 1), 746_180);
|
||||
assert_eq!(tcn_block_params(440, 340, 3, 20, 20, 1), 464_780);
|
||||
assert_eq!(tcn_block_params(340, 240, 3, 20, 20, 1), 249_380);
|
||||
// Conv encoder.
|
||||
assert_eq!(conv_block_params(1, 8), 504);
|
||||
assert_eq!(conv_block_params(8, 8), 728);
|
||||
@@ -347,7 +527,131 @@ mod tests {
|
||||
assert_eq!(conv_block_params(32, 64), 33_472);
|
||||
// Attention + decoder.
|
||||
assert_eq!(axial_attention_params(64, 8), 12_816);
|
||||
assert_eq!(decoder_params(64), 18_598);
|
||||
assert_eq!(decoder_params(64, 32), 18_598);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// ADR-152 efficiency-sweep compact presets. The parameter pins are
|
||||
// GROUND TRUTH measured from the trained PyTorch checkpoints
|
||||
// (benchmarks/wiflow-std/results/efficiency_sweep.jsonl, 2026-06-11):
|
||||
// any mismatch means the Rust formula or config mapping is wrong.
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn half_preset_param_count_matches_trained_checkpoint() {
|
||||
let cfg = WiFlowStdConfig::half();
|
||||
cfg.validate().expect("half preset must validate");
|
||||
assert_eq!(cfg.param_count(), 843_834);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quarter_preset_param_count_matches_trained_checkpoint() {
|
||||
let cfg = WiFlowStdConfig::quarter();
|
||||
cfg.validate().expect("quarter preset must validate");
|
||||
assert_eq!(cfg.param_count(), 338_600);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_preset_param_count_matches_trained_checkpoint() {
|
||||
let cfg = WiFlowStdConfig::tiny();
|
||||
cfg.validate().expect("tiny preset must validate");
|
||||
assert_eq!(cfg.param_count(), 56_290);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preset_tcn_groups_match_sweep_per_block_record() {
|
||||
// efficiency_sweep.jsonl "tcn_groups_per_block": (conv1, conv2) of
|
||||
// each block — conv1 groups over c_in, conv2 over c_out.
|
||||
let half = WiFlowStdConfig::half();
|
||||
let groups: Vec<(usize, usize)> = {
|
||||
let mut c_in = half.subcarriers;
|
||||
half.tcn_channels
|
||||
.iter()
|
||||
.map(|&c_out| {
|
||||
let g = (half.tcn_conv_groups(c_in), half.tcn_conv_groups(c_out));
|
||||
c_in = c_out;
|
||||
g
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
assert_eq!(groups, [(20, 10), (10, 20), (20, 10), (10, 20)]);
|
||||
|
||||
let tiny = WiFlowStdConfig::tiny();
|
||||
assert_eq!(tiny.tcn_conv_groups(540), 540); // depthwise input conv
|
||||
assert_eq!(tiny.tcn_conv_groups(68), 68);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preset_stride_schedules_match_sweep_record() {
|
||||
// efficiency_sweep.jsonl "conv_strides" / "final_width".
|
||||
assert_eq!(WiFlowStdConfig::default().conv_strides(), [2, 2, 2, 2]);
|
||||
assert_eq!(WiFlowStdConfig::half().conv_strides(), [2, 2, 2, 1]);
|
||||
assert_eq!(WiFlowStdConfig::quarter().conv_strides(), [2, 2, 1, 1]);
|
||||
assert_eq!(WiFlowStdConfig::tiny().conv_strides(), [2, 1, 1, 1]);
|
||||
assert_eq!(WiFlowStdConfig::half().feature_width(), 15);
|
||||
assert_eq!(WiFlowStdConfig::quarter().feature_width(), 15);
|
||||
assert_eq!(WiFlowStdConfig::tiny().feature_width(), 16);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fixed_mode_with_defaults_is_unchanged_by_new_knobs() {
|
||||
// The new fields default to upstream behavior: gcd(c, 20) == 20 for
|
||||
// every default channel count, so Gcd mode is also a no-op there.
|
||||
let mut cfg = WiFlowStdConfig::default();
|
||||
assert_eq!(cfg.param_count(), REFERENCE_PARAMS);
|
||||
cfg.tcn_groups_mode = TcnGroupsMode::Gcd;
|
||||
cfg.validate().expect("gcd mode validates at defaults");
|
||||
assert_eq!(cfg.param_count(), REFERENCE_PARAMS);
|
||||
assert_eq!(WiFlowStdConfig::default().decoder_mid(), 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_bad_input_pw_groups() {
|
||||
// 7 divides neither 540 nor 540's first TCN level.
|
||||
let cfg = WiFlowStdConfig {
|
||||
input_pw_groups: 7,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(cfg.validate().is_err());
|
||||
// 27 divides subcarriers=540 but not tiny's tcn_channels[0]=68.
|
||||
let cfg = WiFlowStdConfig {
|
||||
input_pw_groups: 27,
|
||||
..WiFlowStdConfig::tiny()
|
||||
};
|
||||
assert!(cfg.validate().is_err());
|
||||
let zero = WiFlowStdConfig {
|
||||
input_pw_groups: 0,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(zero.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serde_defaults_for_new_fields_are_backward_compatible() {
|
||||
// A config serialized before the compact-variant knobs existed must
|
||||
// deserialize to upstream behavior (Fixed mode, input_pw_groups 1).
|
||||
let legacy = r#"{
|
||||
"subcarriers": 540, "window": 20,
|
||||
"tcn_channels": [540, 440, 340, 240], "tcn_groups": 20,
|
||||
"conv_channels": [8, 16, 32, 64], "attention_groups": 8,
|
||||
"keypoints": 15, "dropout": 0.5
|
||||
}"#;
|
||||
let cfg: WiFlowStdConfig = serde_json::from_str(legacy).expect("deserialize");
|
||||
assert_eq!(cfg, WiFlowStdConfig::default());
|
||||
assert_eq!(cfg.param_count(), REFERENCE_PARAMS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serde_roundtrip_preserves_presets() {
|
||||
for cfg in [
|
||||
WiFlowStdConfig::half(),
|
||||
WiFlowStdConfig::quarter(),
|
||||
WiFlowStdConfig::tiny(),
|
||||
] {
|
||||
let json = serde_json::to_string(&cfg).expect("serialize");
|
||||
let back: WiFlowStdConfig = serde_json::from_str(&json).expect("deserialize");
|
||||
assert_eq!(back, cfg);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -41,12 +41,20 @@ pub(super) struct GroupedTemporalBlock {
|
||||
}
|
||||
|
||||
impl GroupedTemporalBlock {
|
||||
/// `g_in`/`g_out`: group counts of the two grouped convs (each conv
|
||||
/// groups over its own channel count — they differ under the ADR-152
|
||||
/// compact variants' `Gcd`/`Depthwise` modes). `pw_groups` groups the
|
||||
/// first pointwise conv and the residual projection (`input_pw_groups`
|
||||
/// on block 0; 1 everywhere else).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn new(
|
||||
vs: nn::Path,
|
||||
c_in: i64,
|
||||
c_out: i64,
|
||||
dilation: i64,
|
||||
groups: i64,
|
||||
g_in: i64,
|
||||
g_out: i64,
|
||||
pw_groups: i64,
|
||||
dropout: f64,
|
||||
) -> Self {
|
||||
let k = TCN_KERNEL as i64;
|
||||
@@ -58,24 +66,25 @@ impl GroupedTemporalBlock {
|
||||
bias: false,
|
||||
..Default::default()
|
||||
};
|
||||
let pointwise_cfg = nn::ConvConfig {
|
||||
let pointwise_cfg = |groups| nn::ConvConfig {
|
||||
groups,
|
||||
bias: false,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let conv1_group = nn::conv1d(&vs / "conv1_group", c_in, c_in, k, grouped_cfg(groups));
|
||||
let conv1_group = nn::conv1d(&vs / "conv1_group", c_in, c_in, k, grouped_cfg(g_in));
|
||||
let bn1_group = nn::batch_norm1d(&vs / "bn1_group", c_in, bn_cfg());
|
||||
let conv1_pw = nn::conv1d(&vs / "conv1_pw", c_in, c_out, 1, pointwise_cfg);
|
||||
let conv1_pw = nn::conv1d(&vs / "conv1_pw", c_in, c_out, 1, pointwise_cfg(pw_groups));
|
||||
let bn1_pw = nn::batch_norm1d(&vs / "bn1_pw", c_out, bn_cfg());
|
||||
|
||||
let conv2_group = nn::conv1d(&vs / "conv2_group", c_out, c_out, k, grouped_cfg(groups));
|
||||
let conv2_group = nn::conv1d(&vs / "conv2_group", c_out, c_out, k, grouped_cfg(g_out));
|
||||
let bn2_group = nn::batch_norm1d(&vs / "bn2_group", c_out, bn_cfg());
|
||||
let conv2_pw = nn::conv1d(&vs / "conv2_pw", c_out, c_out, 1, pointwise_cfg);
|
||||
let conv2_pw = nn::conv1d(&vs / "conv2_pw", c_out, c_out, 1, pointwise_cfg(1));
|
||||
let bn2_pw = nn::batch_norm1d(&vs / "bn2_pw", c_out, bn_cfg());
|
||||
|
||||
let downsample = (c_in != c_out).then(|| {
|
||||
(
|
||||
nn::conv1d(&vs / "ds_conv", c_in, c_out, 1, pointwise_cfg),
|
||||
nn::conv1d(&vs / "ds_conv", c_in, c_out, 1, pointwise_cfg(pw_groups)),
|
||||
nn::batch_norm1d(&vs / "ds_bn", c_out, bn_cfg()),
|
||||
)
|
||||
});
|
||||
|
||||
@@ -63,7 +63,7 @@ mod layers;
|
||||
#[cfg(feature = "tch-backend")]
|
||||
pub mod model;
|
||||
|
||||
pub use config::WiFlowStdConfig;
|
||||
pub use config::{TcnGroupsMode, WiFlowStdConfig};
|
||||
|
||||
#[cfg(feature = "tch-backend")]
|
||||
pub use model::WiFlowStdModel;
|
||||
|
||||
@@ -59,33 +59,40 @@ impl WiFlowStdModel {
|
||||
let vs = nn::VarStore::new(device);
|
||||
let root = vs.root();
|
||||
|
||||
// TCN stack: dilation doubles per level, causal padding.
|
||||
// TCN stack: dilation doubles per level, causal padding. Per-conv
|
||||
// groups follow `config.tcn_groups_mode`; only block 0's pointwise/
|
||||
// downsample convs use `config.input_pw_groups` (ADR-152 sweep).
|
||||
let mut tcn = Vec::with_capacity(config.tcn_channels.len());
|
||||
let mut c_in = config.subcarriers as i64;
|
||||
let mut c_in = config.subcarriers;
|
||||
for (i, &c_out) in config.tcn_channels.iter().enumerate() {
|
||||
let dilation = 1_i64 << i;
|
||||
let pw_groups = if i == 0 { config.input_pw_groups } else { 1 };
|
||||
tcn.push(GroupedTemporalBlock::new(
|
||||
&root / format!("tcn{i}"),
|
||||
c_in,
|
||||
c_in as i64,
|
||||
c_out as i64,
|
||||
dilation,
|
||||
config.tcn_groups as i64,
|
||||
config.tcn_conv_groups(c_in) as i64,
|
||||
config.tcn_conv_groups(c_out) as i64,
|
||||
pw_groups as i64,
|
||||
config.dropout,
|
||||
));
|
||||
c_in = c_out as i64;
|
||||
c_in = c_out;
|
||||
}
|
||||
|
||||
// 2-D conv encoder: ConvBlock1 (stride 1) + strided asymmetric blocks.
|
||||
// 2-D conv encoder: ConvBlock1 (stride 1) + asymmetric blocks with
|
||||
// the derived stride schedule ([2, 2, 2, 2] at the upstream default).
|
||||
let c0 = config.conv_channels[0] as i64;
|
||||
let conv_in = ConvBlock::new(&root / "conv_in", 1, c0, 1);
|
||||
let mut conv_blocks = Vec::with_capacity(config.conv_channels.len());
|
||||
let strides = config.conv_strides();
|
||||
let mut c_in = c0;
|
||||
for (i, &c_out) in config.conv_channels.iter().enumerate() {
|
||||
conv_blocks.push(ConvBlock::new(
|
||||
&root / format!("conv{i}"),
|
||||
c_in,
|
||||
c_out as i64,
|
||||
2,
|
||||
strides[i] as i64,
|
||||
));
|
||||
c_in = c_out as i64;
|
||||
}
|
||||
@@ -93,8 +100,8 @@ impl WiFlowStdModel {
|
||||
let attention =
|
||||
DualAxialAttention::new(&root / "attention", c_in, config.attention_groups as i64);
|
||||
|
||||
// Decoder: c → c/2 (3×3) → 2 (1×1), BN + SiLU after each conv.
|
||||
let mid = c_in / 2;
|
||||
// Decoder: c → decoder_mid (3×3) → 2 (1×1), BN + SiLU after each conv.
|
||||
let mid = config.decoder_mid() as i64;
|
||||
let dec_conv1 = nn::conv2d(
|
||||
&root / "dec_conv1",
|
||||
c_in,
|
||||
@@ -241,6 +248,26 @@ mod tests {
|
||||
assert_eq!(model.num_parameters(), 2_225_042);
|
||||
}
|
||||
|
||||
/// ADR-152 efficiency-sweep compact presets: the tch graph must realise
|
||||
/// exactly the trained checkpoints' measured parameter counts
|
||||
/// (benchmarks/wiflow-std/results/efficiency_sweep.jsonl) and produce
|
||||
/// the standard [B, 15, 2] output.
|
||||
#[test]
|
||||
fn compact_preset_param_counts_and_shapes() {
|
||||
for (cfg, expected) in [
|
||||
(WiFlowStdConfig::half(), 843_834_i64),
|
||||
(WiFlowStdConfig::quarter(), 338_600),
|
||||
(WiFlowStdConfig::tiny(), 56_290),
|
||||
] {
|
||||
tch::manual_seed(0);
|
||||
let model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("preset builds");
|
||||
assert_eq!(model.num_parameters(), expected);
|
||||
assert_eq!(model.num_parameters(), cfg.param_count() as i64);
|
||||
let out = model.forward_inference(&random_csi(&cfg, 2));
|
||||
assert_eq!(out.size(), &[2, 15, 2]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forward_output_shape_15_keypoints() {
|
||||
tch::manual_seed(0);
|
||||
|
||||
Reference in New Issue
Block a user