From af6621005f4517e90066d3f4c54084e81f671b29 Mon Sep 17 00:00:00 2001 From: ruv Date: Wed, 10 Jun 2026 20:23:38 -0400 Subject: [PATCH] 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 --- benchmarks/wiflow-std/.gitignore | 2 + .../wiflow-std/export_to_safetensors.py | 187 ++++++++++++++++++ .../src/wiflow_std/model.rs | 27 ++- .../tests/test_wiflow_std_parity.rs | 93 +++++++++ 4 files changed, 307 insertions(+), 2 deletions(-) create mode 100644 benchmarks/wiflow-std/export_to_safetensors.py create mode 100644 v2/crates/wifi-densepose-train/tests/test_wiflow_std_parity.rs diff --git a/benchmarks/wiflow-std/.gitignore b/benchmarks/wiflow-std/.gitignore index 772dcbf8..6d608ac1 100644 --- a/benchmarks/wiflow-std/.gitignore +++ b/benchmarks/wiflow-std/.gitignore @@ -13,4 +13,6 @@ downloads/ *.npz *.zip *.mat +*.safetensors +results/parity_fixture.json __pycache__/ diff --git a/benchmarks/wiflow-std/export_to_safetensors.py b/benchmarks/wiflow-std/export_to_safetensors.py new file mode 100644 index 00000000..b5c9d4d5 --- /dev/null +++ b/benchmarks/wiflow-std/export_to_safetensors.py @@ -0,0 +1,187 @@ +"""ADR-152 §2.2: export the retrained WiFlow-STD PyTorch checkpoint to +safetensors with tch-rs (VarStore) variable names, plus a numerical-parity +fixture for the Rust port. + +Outputs (all under results/, gitignored): + retrained_wiflow_std.safetensors -- 248 f32 tensors named exactly as the + Rust WiFlowStdModel VarStore expects + (see wiflow_std/model.rs + `dump_variable_names` for the + authoritative name dump) + parity_fixture.npz -- deterministic input (seed 42, + shape (2, 540, 20), uniform [0,1]) and + the Python model's eval-mode output + parity_fixture.json -- same data as flattened f32 lists, for + the dependency-free Rust test + (tests/test_wiflow_std_parity.rs) + +PyTorch -> tch key mapping (derived from the VarStore dump, not guessed): + + tcn.network.{i}.conv1_group.weight -> tcn{i}.conv1_group.weight + tcn.network.{i}.bn*_{group,pw}. -> tcn{i}.bn*_{group,pw}. + tcn.network.{i}.downsample.0.weight -> tcn{i}.ds_conv.weight + tcn.network.{i}.downsample.1. -> tcn{i}.ds_bn. + up.block.{0,1,4,5,8,9}. -> conv_in.{conv1,bn1,conv2,bn2,conv3,bn3}. + up.downsample.{0,1}. -> conv_in.{ds_conv,ds_bn}. + residual_blocks.{i}.block.{...}. -> conv{i}.{conv1..bn3}. + residual_blocks.{i}.downsample.{0,1} -> conv{i}.{ds_conv,ds_bn} + attention.{width,height}_axis.qkv_transform.weight + -> attention.{width,height}.qkv.weight + attention.{width,height}_axis.bn_* -> attention.{width,height}.bn_* + decoder.{0,1,3,4}. -> {dec_conv1,dec_bn1,dec_conv2,dec_bn2}. + *.num_batches_tracked -> dropped (tch BatchNorm has no such buffer) + +Legacy upstream names (att. -> attention., final_conv. -> decoder.) are +remapped first, exactly as eval_repro.py does for the released checkpoint. + +Usage: + .venv/Scripts/python.exe export_to_safetensors.py +""" + +import json +import os +import re +import sys + +import numpy as np +import torch +from safetensors.torch import save_file + +HERE = os.path.dirname(os.path.abspath(__file__)) +UPSTREAM = os.path.join(HERE, "upstream") +RESULTS = os.path.join(HERE, "results") +sys.path.insert(0, UPSTREAM) + +# Upstream models/__init__.py is broken as published (imports a name tcn.py +# does not define); register a stub package so it never executes. +import types # noqa: E402 + +_models_pkg = types.ModuleType("models") +_models_pkg.__path__ = [os.path.join(UPSTREAM, "models")] +sys.modules["models"] = _models_pkg + +from models.pose_model import WiFlowPoseModel # noqa: E402 + +CHECKPOINT = os.path.join(RESULTS, "retrained_best_pose_model.pth") + +# Sequential index -> tch sub-name inside one ConvBlock1/AsymmetricConvBlock: +# [Conv2d(0), BN(1), SiLU(2), Dropout2d(3), Conv2d(4), BN(5), SiLU(6), +# Dropout2d(7), Conv2d(8), BN(9)] +_BLOCK_IDX = {"0": "conv1", "1": "bn1", "4": "conv2", "5": "bn2", + "8": "conv3", "9": "bn3"} +_DS_IDX = {"0": "ds_conv", "1": "ds_bn"} +_DECODER_IDX = {"0": "dec_conv1", "1": "dec_bn1", "3": "dec_conv2", + "4": "dec_bn2"} + + +def _conv_block(new_prefix: str, rest: str) -> str: + m = re.fullmatch(r"block\.(\d+)\.(.+)", rest) + if m: + return f"{new_prefix}.{_BLOCK_IDX[m.group(1)]}.{m.group(2)}" + m = re.fullmatch(r"downsample\.(\d+)\.(.+)", rest) + if m: + return f"{new_prefix}.{_DS_IDX[m.group(1)]}.{m.group(2)}" + raise KeyError(f"unmapped conv-block key: {new_prefix} / {rest}") + + +def map_key(key: str) -> str: + """Map one PyTorch state_dict key to the tch VarStore name.""" + m = re.fullmatch(r"tcn\.network\.(\d+)\.(.+)", key) + if m: + i, rest = m.groups() + rest = (rest.replace("downsample.0.", "ds_conv.") + .replace("downsample.1.", "ds_bn.")) + return f"tcn{i}.{rest}" + + m = re.fullmatch(r"up\.(.+)", key) + if m: + return _conv_block("conv_in", m.group(1)) + + m = re.fullmatch(r"residual_blocks\.(\d+)\.(.+)", key) + if m: + return _conv_block(f"conv{m.group(1)}", m.group(2)) + + m = re.fullmatch(r"attention\.(width|height)_axis\.(.+)", key) + if m: + axis, rest = m.groups() + rest = rest.replace("qkv_transform.", "qkv.") + return f"attention.{axis}.{rest}" + + m = re.fullmatch(r"decoder\.(\d+)\.(.+)", key) + if m: + return f"{_DECODER_IDX[m.group(1)]}.{m.group(2)}" + + raise KeyError(f"unmapped checkpoint key: {key}") + + +def main(): + state = torch.load(CHECKPOINT, map_location="cpu", weights_only=True) + if not isinstance(state, dict) or "tcn.network.0.conv1_group.weight" not in { + k for k in state + } | {k.replace("att.", "attention.") for k in state}: + # tolerate trainer wrappers like {"model_state_dict": ...} + for wrapper in ("model_state_dict", "state_dict", "model"): + if isinstance(state, dict) and wrapper in state: + state = state[wrapper] + break + + # Legacy upstream names predate the published code (eval_repro.py). + renames = {"att.": "attention.", "final_conv.": "decoder."} + state = {next((new + k[len(old):] for old, new in renames.items() + if k.startswith(old)), k): v + for k, v in state.items()} + + mapped = {} + dropped = 0 + for k, v in state.items(): + if k.endswith("num_batches_tracked"): + dropped += 1 + continue + tch_key = map_key(k) + if tch_key in mapped: + raise KeyError(f"duplicate mapped key: {k} -> {tch_key}") + mapped[tch_key] = v.detach().to(torch.float32).contiguous() + + n_params = sum(v.numel() for k, v in mapped.items() + if "running_" not in k) + print(f"checkpoint tensors: {len(state)} " + f"(dropped {dropped} num_batches_tracked)") + print(f"mapped tensors: {len(mapped)}, " + f"non-buffer params: {n_params/1e6:.6f}M") + assert len(mapped) == 248, f"expected 248 tch variables, got {len(mapped)}" + assert n_params == 2_225_042, f"param count mismatch: {n_params}" + + st_path = os.path.join(RESULTS, "retrained_wiflow_std.safetensors") + save_file(mapped, st_path) + print(f"wrote {st_path}") + + # ---- parity fixture -------------------------------------------------- + model = WiFlowPoseModel(dropout=0.5) + model.load_state_dict(state, strict=True) + model.eval() + + gen = torch.Generator().manual_seed(42) + x = torch.rand(2, 540, 20, generator=gen, dtype=torch.float32) + with torch.no_grad(): + y = model(x) + print(f"fixture input {tuple(x.shape)} -> output {tuple(y.shape)}, " + f"output range [{y.min().item():.6f}, {y.max().item():.6f}]") + + np.savez(os.path.join(RESULTS, "parity_fixture.npz"), + input=x.numpy(), output=y.numpy()) + fixture = { + "seed": 42, + "input_shape": list(x.shape), + "input": x.flatten().tolist(), + "output_shape": list(y.shape), + "output": y.flatten().tolist(), + } + json_path = os.path.join(RESULTS, "parity_fixture.json") + with open(json_path, "w") as f: + json.dump(fixture, f) + print(f"wrote {os.path.join(RESULTS, 'parity_fixture.npz')}") + print(f"wrote {json_path}") + + +if __name__ == "__main__": + main() diff --git a/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs b/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs index fa408e74..0ede0632 100644 --- a/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs +++ b/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs @@ -3,8 +3,12 @@ //! 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`). +//! qkv conv mirrors the reference's `N(0, sqrt(1/in_planes))` init). 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}; @@ -267,6 +271,25 @@ mod tests { ); } + /// 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)> = + 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 { diff --git a/v2/crates/wifi-densepose-train/tests/test_wiflow_std_parity.rs b/v2/crates/wifi-densepose-train/tests/test_wiflow_std_parity.rs new file mode 100644 index 00000000..4b19ff7c --- /dev/null +++ b/v2/crates/wifi-densepose-train/tests/test_wiflow_std_parity.rs @@ -0,0 +1,93 @@ +//! Numerical parity between the Rust WiFlow-STD port and the retrained +//! PyTorch checkpoint (ADR-152 §2.2). +//! +//! The fixtures are produced by `benchmarks/wiflow-std/export_to_safetensors.py` +//! (gitignored — they derive from the retrained checkpoint, which is itself +//! gitignored): +//! +//! - `results/retrained_wiflow_std.safetensors` — the epoch-36 checkpoint +//! (val PCK@20 96.99%) remapped to tch `VarStore` variable names +//! - `results/parity_fixture.json` — a deterministic input (seed 42, shape +//! `(2, 540, 20)`, uniform `[0, 1]`) and the upstream `WiFlowPoseModel`'s +//! eval-mode output on it +//! +//! Run explicitly (needs LibTorch, e.g. `LIBTORCH_USE_PYTORCH=1` with the +//! torch DLL directory on `PATH`): +//! +//! ```text +//! cargo test -p wifi-densepose-train --features tch-backend \ +//! --test test_wiflow_std_parity -- --ignored --nocapture +//! ``` + +#![cfg(feature = "tch-backend")] + +use std::fs::File; +use std::io::BufReader; +use std::path::PathBuf; + +use tch::{Device, Tensor}; +use wifi_densepose_train::{WiFlowStdConfig, WiFlowStdModel}; + +#[derive(serde::Deserialize)] +struct ParityFixture { + input_shape: Vec, + input: Vec, + output_shape: Vec, + output: Vec, +} + +fn results_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("..") + .join("..") + .join("..") + .join("benchmarks") + .join("wiflow-std") + .join("results") +} + +/// Loads the retrained checkpoint into the Rust model and asserts the forward +/// pass matches PyTorch to within 1e-4 max absolute difference. +/// +/// `#[ignore]`d by default: it needs the gitignored fixtures above plus a +/// working LibTorch environment, neither of which exist in CI. +#[test] +#[ignore = "needs gitignored fixtures (run export_to_safetensors.py) + LibTorch env; run with --ignored"] +fn retrained_checkpoint_matches_pytorch_forward() { + let dir = results_dir(); + let weights = dir.join("retrained_wiflow_std.safetensors"); + let fixture_path = dir.join("parity_fixture.json"); + for p in [&weights, &fixture_path] { + assert!( + p.exists(), + "missing fixture {} — run benchmarks/wiflow-std/export_to_safetensors.py first", + p.display() + ); + } + + let fixture: ParityFixture = serde_json::from_reader(BufReader::new( + File::open(&fixture_path).expect("open parity_fixture.json"), + )) + .expect("parse parity_fixture.json"); + assert_eq!(fixture.input_shape, vec![2, 540, 20]); + assert_eq!(fixture.output_shape, vec![2, 15, 2]); + + let cfg = WiFlowStdConfig::default(); + let mut model = WiFlowStdModel::new(&cfg, Device::Cpu).expect("build default model"); + model + .load(&weights) + .expect("safetensors load: every VarStore variable must match by name and shape"); + + let input = Tensor::from_slice(&fixture.input).reshape(&fixture.input_shape[..]); + let expected = Tensor::from_slice(&fixture.output).reshape(&fixture.output_shape[..]); + + let output = model.forward_inference(&input); + assert_eq!(output.size(), fixture.output_shape); + + let max_diff = (&output - &expected).abs().max().double_value(&[]); + println!("max |rust - python| = {max_diff:.3e}"); + assert!( + max_diff < 1e-4, + "Rust forward pass diverges from PyTorch: max abs diff {max_diff:.3e} >= 1e-4" + ); +}