dearsky/wifi-densepose
1
0
Fork 0
Code Issues 20 Pull Requests 5 Actions Packages Projects Releases 1 Wiki Activity

feat: ADR-024 AETHER — Contrastive CSI Embedding Model

Browse Source 
Implements Project AETHER (Ambient Electromagnetic Topology for
Hierarchical Embedding and Recognition): self-supervised contrastive
learning for WiFi CSI fingerprinting, similarity search, and anomaly
detection.

New files:
- docs/adr/ADR-024 — full architectural spec (1024 lines) with
  mathematical foundations, 6 implementation phases, 30 SOTA references
- embedding.rs — ProjectionHead, CsiAugmenter, InfoNCE loss,
  FingerprintIndex, PoseEncoder, EmbeddingExtractor (909 lines)

Modified:
- main.rs — CLI flags: --pretrain, --pretrain-epochs, --embed, --build-index
- trainer.rs — contrastive pretraining loop integration
- graph_transformer.rs — body_part_features exposure for embedding extraction
- rvf_container.rs — embedding segment type support
- lib.rs — embedding module export
- README.md — collapsible AETHER section with architecture, training modes,
  index types, and model size table

53K params total, fits in 55 KB on ESP32. No external ML dependencies.

Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruv
2026-03-01 01:18:30 -05:00
parent 44b9c30dbc
commit 5942d4dd5b
8 changed files with 2526 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

909
rust-port/wifi-densepose-rs/crates/wifi-densepose-sensing-server/src/embedding.rs Normal file
View File

@@ -0,0 +1,909 @@
//! Contrastive CSI Embedding Model (ADR-024).
//!
//! Implements self-supervised contrastive learning for WiFi CSI feature extraction:
//! - ProjectionHead: 2-layer MLP for contrastive embedding space
//! - CsiAugmenter: domain-specific augmentations for SimCLR-style pretraining
//! - InfoNCE loss: normalized temperature-scaled cross-entropy
//! - FingerprintIndex: brute-force nearest-neighbour (HNSW-compatible interface)
//! - PoseEncoder: lightweight encoder for cross-modal alignment
//! - EmbeddingExtractor: full pipeline (backbone + projection)
//!
//! All arithmetic uses `f32`. No external ML dependencies.
use crate::graph_transformer::{CsiToPoseTransformer, TransformerConfig, Linear};
// ── SimpleRng (xorshift64) ──────────────────────────────────────────────────
/// Deterministic xorshift64 PRNG to avoid external dependency.
struct SimpleRng {
state: u64,
}
impl SimpleRng {
fn new(seed: u64) -> Self {
Self { state: if seed == 0 { 0xBAAD_CAFE_DEAD_BEEFu64 } else { seed } }
}
fn next_u64(&mut self) -> u64 {
let mut x = self.state;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
self.state = x;
x
}
/// Uniform f32 in [0, 1).
fn next_f32_unit(&mut self) -> f32 {
(self.next_u64() >> 11) as f32 / (1u64 << 53) as f32
}
/// Gaussian approximation via Box-Muller (pair, returns first).
fn next_gaussian(&mut self) -> f32 {
let u1 = self.next_f32_unit().max(1e-10);
let u2 = self.next_f32_unit();
(-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos()
}
}
// ── EmbeddingConfig ─────────────────────────────────────────────────────────
/// Configuration for the contrastive embedding model.
#[derive(Debug, Clone)]
pub struct EmbeddingConfig {
/// Hidden dimension (must match transformer d_model).
pub d_model: usize,
/// Projection/embedding dimension.
pub d_proj: usize,
/// InfoNCE temperature.
pub temperature: f32,
/// Whether to L2-normalize output embeddings.
pub normalize: bool,
}
impl Default for EmbeddingConfig {
fn default() -> Self {
Self { d_model: 64, d_proj: 128, temperature: 0.07, normalize: true }
}
}
// ── ProjectionHead ──────────────────────────────────────────────────────────
/// 2-layer MLP projection head: d_model -> d_proj -> d_proj with ReLU + L2-norm.
#[derive(Debug, Clone)]
pub struct ProjectionHead {
pub proj_1: Linear,
pub proj_2: Linear,
pub config: EmbeddingConfig,
}
impl ProjectionHead {
/// Xavier-initialized projection head.
pub fn new(config: EmbeddingConfig) -> Self {
Self {
proj_1: Linear::with_seed(config.d_model, config.d_proj, 2024),
proj_2: Linear::with_seed(config.d_proj, config.d_proj, 2025),
config,
}
}
/// Zero-initialized projection head (for gradient estimation).
pub fn zeros(config: EmbeddingConfig) -> Self {
Self {
proj_1: Linear::zeros(config.d_model, config.d_proj),
proj_2: Linear::zeros(config.d_proj, config.d_proj),
config,
}
}
/// Forward pass: ReLU between layers, optional L2-normalize output.
pub fn forward(&self, x: &[f32]) -> Vec<f32> {
let h: Vec<f32> = self.proj_1.forward(x).into_iter()
.map(|v| if v > 0.0 { v } else { 0.0 })
.collect();
let mut out = self.proj_2.forward(&h);
if self.config.normalize {
l2_normalize(&mut out);
}
out
}
/// Push all weights into a flat vec.
pub fn flatten_into(&self, out: &mut Vec<f32>) {
self.proj_1.flatten_into(out);
self.proj_2.flatten_into(out);
}
/// Restore from a flat slice. Returns (Self, number of f32s consumed).
pub fn unflatten_from(data: &[f32], config: &EmbeddingConfig) -> (Self, usize) {
let mut offset = 0;
let (p1, n) = Linear::unflatten_from(&data[offset..], config.d_model, config.d_proj);
offset += n;
let (p2, n) = Linear::unflatten_from(&data[offset..], config.d_proj, config.d_proj);
offset += n;
(Self { proj_1: p1, proj_2: p2, config: config.clone() }, offset)
}
/// Total trainable parameters.
pub fn param_count(&self) -> usize {
self.proj_1.param_count() + self.proj_2.param_count()
}
}
// ── CsiAugmenter ────────────────────────────────────────────────────────────
/// CSI augmentation strategies for contrastive pretraining.
#[derive(Debug, Clone)]
pub struct CsiAugmenter {
/// +/- frames to shift (temporal jitter).
pub temporal_jitter: i32,
/// Fraction of subcarriers to zero out.
pub subcarrier_mask_ratio: f32,
/// Gaussian noise sigma.
pub noise_std: f32,
/// Max phase offset in radians.
pub phase_rotation_max: f32,
/// Amplitude scale range (min, max).
pub amplitude_scale_range: (f32, f32),
}
impl CsiAugmenter {
pub fn new() -> Self {
Self {
temporal_jitter: 2,
subcarrier_mask_ratio: 0.15,
noise_std: 0.05,
phase_rotation_max: std::f32::consts::FRAC_PI_4,
amplitude_scale_range: (0.8, 1.2),
}
}
/// Apply random augmentations to a CSI window, returning two different views.
/// Each view receives a different random subset of augmentations.
pub fn augment_pair(
&self,
csi_window: &[Vec<f32>],
rng_seed: u64,
) -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
let mut rng_a = SimpleRng::new(rng_seed);
let mut rng_b = SimpleRng::new(rng_seed.wrapping_add(0x1234_5678_9ABC_DEF0));
// View A: temporal jitter + noise + subcarrier mask
let mut view_a = self.apply_temporal_jitter(csi_window, &mut rng_a);
self.apply_gaussian_noise(&mut view_a, &mut rng_a);
self.apply_subcarrier_mask(&mut view_a, &mut rng_a);
// View B: amplitude scaling + phase rotation + different noise
let mut view_b = self.apply_temporal_jitter(csi_window, &mut rng_b);
self.apply_amplitude_scaling(&mut view_b, &mut rng_b);
self.apply_phase_rotation(&mut view_b, &mut rng_b);
self.apply_gaussian_noise(&mut view_b, &mut rng_b);
(view_a, view_b)
}
fn apply_temporal_jitter(
&self,
window: &[Vec<f32>],
rng: &mut SimpleRng,
) -> Vec<Vec<f32>> {
if window.is_empty() || self.temporal_jitter == 0 {
return window.to_vec();
}
let range = 2 * self.temporal_jitter + 1;
let shift = (rng.next_u64() % range as u64) as i32 - self.temporal_jitter;
let n = window.len() as i32;
(0..window.len())
.map(|i| {
let src = (i as i32 + shift).clamp(0, n - 1) as usize;
window[src].clone()
})
.collect()
}
fn apply_subcarrier_mask(&self, window: &mut [Vec<f32>], rng: &mut SimpleRng) {
for frame in window.iter_mut() {
for v in frame.iter_mut() {
if rng.next_f32_unit() < self.subcarrier_mask_ratio {
*v = 0.0;
}
}
}
}
fn apply_gaussian_noise(&self, window: &mut [Vec<f32>], rng: &mut SimpleRng) {
for frame in window.iter_mut() {
for v in frame.iter_mut() {
*v += rng.next_gaussian() * self.noise_std;
}
}
}
fn apply_phase_rotation(&self, window: &mut [Vec<f32>], rng: &mut SimpleRng) {
let offset = (rng.next_f32_unit() * 2.0 - 1.0) * self.phase_rotation_max;
for frame in window.iter_mut() {
for v in frame.iter_mut() {
// Approximate phase rotation on amplitude: multiply by cos(offset)
*v *= offset.cos();
}
}
}
fn apply_amplitude_scaling(&self, window: &mut [Vec<f32>], rng: &mut SimpleRng) {
let (lo, hi) = self.amplitude_scale_range;
let scale = lo + rng.next_f32_unit() * (hi - lo);
for frame in window.iter_mut() {
for v in frame.iter_mut() {
*v *= scale;
}
}
}
}
impl Default for CsiAugmenter {
fn default() -> Self { Self::new() }
}
// ── Vector math utilities ───────────────────────────────────────────────────
/// L2-normalize a vector in-place.
fn l2_normalize(v: &mut [f32]) {
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 1e-10 {
let inv = 1.0 / norm;
for x in v.iter_mut() {
*x *= inv;
}
}
}
/// Cosine similarity between two vectors.
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
let n = a.len().min(b.len());
let dot: f32 = (0..n).map(|i| a[i] * b[i]).sum();
let na = (0..n).map(|i| a[i] * a[i]).sum::<f32>().sqrt();
let nb = (0..n).map(|i| b[i] * b[i]).sum::<f32>().sqrt();
if na > 1e-10 && nb > 1e-10 { dot / (na * nb) } else { 0.0 }
}
// ── InfoNCE loss ────────────────────────────────────────────────────────────
/// InfoNCE contrastive loss (NT-Xent / SimCLR objective).
///
/// For batch of N pairs (a_i, b_i):
/// loss = -1/N sum_i log( exp(sim(a_i, b_i)/t) / sum_j exp(sim(a_i, b_j)/t) )
pub fn info_nce_loss(
embeddings_a: &[Vec<f32>],
embeddings_b: &[Vec<f32>],
temperature: f32,
) -> f32 {
let n = embeddings_a.len().min(embeddings_b.len());
if n == 0 {
return 0.0;
}
let t = temperature.max(1e-6);
let mut total_loss = 0.0f32;
for i in 0..n {
// Compute similarity of anchor a_i with all b_j
let mut logits = Vec::with_capacity(n);
for j in 0..n {
logits.push(cosine_similarity(&embeddings_a[i], &embeddings_b[j]) / t);
}
// Numerically stable log-softmax
let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
let log_sum_exp = logits.iter()
.map(|&l| (l - max_logit).exp())
.sum::<f32>()
.ln() + max_logit;
total_loss += -logits[i] + log_sum_exp;
}
total_loss / n as f32
}
// ── FingerprintIndex ────────────────────────────────────────────────────────
/// Fingerprint index type.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum IndexType {
EnvironmentFingerprint,
ActivityPattern,
TemporalBaseline,
PersonTrack,
}
/// A single index entry.
pub struct IndexEntry {
pub embedding: Vec<f32>,
pub metadata: String,
pub timestamp_ms: u64,
pub index_type: IndexType,
}
/// Search result from the fingerprint index.
pub struct SearchResult {
/// Index into the entries vec.
pub entry: usize,
/// Cosine distance (1 - similarity).
pub distance: f32,
/// Metadata string from the matching entry.
pub metadata: String,
}
/// Brute-force fingerprint index with HNSW-compatible interface.
///
/// Stores embeddings and supports nearest-neighbour search via cosine distance.
/// Can be replaced with a proper HNSW implementation for production scale.
pub struct FingerprintIndex {
entries: Vec<IndexEntry>,
index_type: IndexType,
}
impl FingerprintIndex {
pub fn new(index_type: IndexType) -> Self {
Self { entries: Vec::new(), index_type }
}
/// Insert an embedding with metadata and timestamp.
pub fn insert(&mut self, embedding: Vec<f32>, metadata: String, timestamp_ms: u64) {
self.entries.push(IndexEntry {
embedding,
metadata,
timestamp_ms,
index_type: self.index_type,
});
}
/// Search for the top-k nearest embeddings by cosine distance.
pub fn search(&self, query: &[f32], top_k: usize) -> Vec<SearchResult> {
let mut results: Vec<(usize, f32)> = self.entries.iter().enumerate()
.map(|(i, e)| (i, 1.0 - cosine_similarity(query, &e.embedding)))
.collect();
results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
results.truncate(top_k);
results.into_iter().map(|(i, d)| SearchResult {
entry: i,
distance: d,
metadata: self.entries[i].metadata.clone(),
}).collect()
}
/// Number of entries in the index.
pub fn len(&self) -> usize { self.entries.len() }
/// Whether the index is empty.
pub fn is_empty(&self) -> bool { self.entries.is_empty() }
/// Detect anomaly: returns true if query is farther than threshold from all entries.
pub fn is_anomaly(&self, query: &[f32], threshold: f32) -> bool {
if self.entries.is_empty() {
return true;
}
self.entries.iter()
.all(|e| (1.0 - cosine_similarity(query, &e.embedding)) > threshold)
}
}
// ── PoseEncoder (cross-modal alignment) ─────────────────────────────────────
/// Lightweight pose encoder for cross-modal alignment.
/// Maps 51-dim pose vector (17 keypoints * 3 coords) to d_proj embedding.
#[derive(Debug, Clone)]
pub struct PoseEncoder {
pub layer_1: Linear,
pub layer_2: Linear,
d_proj: usize,
}
impl PoseEncoder {
/// Create a new pose encoder mapping 51-dim input to d_proj-dim embedding.
pub fn new(d_proj: usize) -> Self {
Self {
layer_1: Linear::with_seed(51, d_proj, 3001),
layer_2: Linear::with_seed(d_proj, d_proj, 3002),
d_proj,
}
}
/// Forward pass: ReLU + L2-normalize.
pub fn forward(&self, pose_flat: &[f32]) -> Vec<f32> {
let h: Vec<f32> = self.layer_1.forward(pose_flat).into_iter()
.map(|v| if v > 0.0 { v } else { 0.0 })
.collect();
let mut out = self.layer_2.forward(&h);
l2_normalize(&mut out);
out
}
/// Push all weights into a flat vec.
pub fn flatten_into(&self, out: &mut Vec<f32>) {
self.layer_1.flatten_into(out);
self.layer_2.flatten_into(out);
}
/// Restore from a flat slice. Returns (Self, number of f32s consumed).
pub fn unflatten_from(data: &[f32], d_proj: usize) -> (Self, usize) {
let mut offset = 0;
let (l1, n) = Linear::unflatten_from(&data[offset..], 51, d_proj);
offset += n;
let (l2, n) = Linear::unflatten_from(&data[offset..], d_proj, d_proj);
offset += n;
(Self { layer_1: l1, layer_2: l2, d_proj }, offset)
}
/// Total trainable parameters.
pub fn param_count(&self) -> usize {
self.layer_1.param_count() + self.layer_2.param_count()
}
}
/// Cross-modal contrastive loss: aligns CSI embeddings with pose embeddings.
/// Same as info_nce_loss but between two different modalities.
pub fn cross_modal_loss(
csi_embeddings: &[Vec<f32>],
pose_embeddings: &[Vec<f32>],
temperature: f32,
) -> f32 {
info_nce_loss(csi_embeddings, pose_embeddings, temperature)
}
// ── EmbeddingExtractor ──────────────────────────────────────────────────────
/// Full embedding extractor: CsiToPoseTransformer backbone + ProjectionHead.
pub struct EmbeddingExtractor {
pub transformer: CsiToPoseTransformer,
pub projection: ProjectionHead,
pub config: EmbeddingConfig,
}
impl EmbeddingExtractor {
/// Create a new embedding extractor with given configs.
pub fn new(t_config: TransformerConfig, e_config: EmbeddingConfig) -> Self {
Self {
transformer: CsiToPoseTransformer::new(t_config),
projection: ProjectionHead::new(e_config.clone()),
config: e_config,
}
}
/// Extract embedding from CSI features.
/// Mean-pools the 17 body_part_features from the transformer backbone,
/// then projects through the ProjectionHead.
pub fn extract(&self, csi_features: &[Vec<f32>]) -> Vec<f32> {
let body_feats = self.transformer.embed(csi_features);
let d = self.config.d_model;
// Mean-pool across 17 keypoints
let mut pooled = vec![0.0f32; d];
for feat in &body_feats {
for (p, &f) in pooled.iter_mut().zip(feat.iter()) {
*p += f;
}
}
let n = body_feats.len() as f32;
if n > 0.0 {
for p in pooled.iter_mut() {
*p /= n;
}
}
self.projection.forward(&pooled)
}
/// Batch extract embeddings.
pub fn extract_batch(&self, batch: &[Vec<Vec<f32>>]) -> Vec<Vec<f32>> {
batch.iter().map(|csi| self.extract(csi)).collect()
}
/// Total parameter count (transformer + projection).
pub fn param_count(&self) -> usize {
self.transformer.param_count() + self.projection.param_count()
}
/// Flatten all weights (transformer + projection).
pub fn flatten_weights(&self) -> Vec<f32> {
let mut out = self.transformer.flatten_weights();
self.projection.flatten_into(&mut out);
out
}
/// Unflatten all weights from a flat slice.
pub fn unflatten_weights(&mut self, params: &[f32]) -> Result<(), String> {
let t_count = self.transformer.param_count();
let p_count = self.projection.param_count();
let expected = t_count + p_count;
if params.len() != expected {
return Err(format!(
"expected {} params ({}+{}), got {}",
expected, t_count, p_count, params.len()
));
}
self.transformer.unflatten_weights(&params[..t_count])?;
let (proj, consumed) = ProjectionHead::unflatten_from(&params[t_count..], &self.config);
if consumed != p_count {
return Err(format!(
"projection consumed {consumed} params, expected {p_count}"
));
}
self.projection = proj;
Ok(())
}
}
// ── Quantized embedding validation ─────────────────────────────────────────
use crate::sparse_inference::Quantizer;
/// Validate that INT8 quantization preserves embedding ranking.
/// Returns Spearman rank correlation between FP32 and INT8 distance rankings.
pub fn validate_quantized_embeddings(
embeddings_fp32: &[Vec<f32>],
query_fp32: &[f32],
_quantizer: &Quantizer,
) -> f32 {
if embeddings_fp32.is_empty() {
return 1.0;
}
let n = embeddings_fp32.len();
// 1. FP32 cosine distances
let fp32_distances: Vec<f32> = embeddings_fp32.iter()
.map(|e| 1.0 - cosine_similarity(query_fp32, e))
.collect();
// 2. Quantize each embedding and query, compute approximate distances
let query_quant = Quantizer::quantize_symmetric(query_fp32);
let query_deq = Quantizer::dequantize(&query_quant);
let int8_distances: Vec<f32> = embeddings_fp32.iter()
.map(|e| {
let eq = Quantizer::quantize_symmetric(e);
let ed = Quantizer::dequantize(&eq);
1.0 - cosine_similarity(&query_deq, &ed)
})
.collect();
// 3. Compute rank arrays
let fp32_ranks = rank_array(&fp32_distances);
let int8_ranks = rank_array(&int8_distances);
// 4. Spearman rank correlation: 1 - 6*sum(d^2) / (n*(n^2-1))
let d_sq_sum: f32 = fp32_ranks.iter().zip(int8_ranks.iter())
.map(|(&a, &b)| (a - b) * (a - b))
.sum();
let n_f = n as f32;
if n <= 1 {
return 1.0;
}
1.0 - (6.0 * d_sq_sum) / (n_f * (n_f * n_f - 1.0))
}
/// Compute ranks for an array of values (1-based, average ties).
fn rank_array(values: &[f32]) -> Vec<f32> {
let n = values.len();
let mut indexed: Vec<(usize, f32)> = values.iter().copied().enumerate().collect();
indexed.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
let mut ranks = vec![0.0f32; n];
let mut i = 0;
while i < n {
let mut j = i;
while j < n && (indexed[j].1 - indexed[i].1).abs() < 1e-10 {
j += 1;
}
let avg_rank = (i + j + 1) as f32 / 2.0; // 1-based average
for k in i..j {
ranks[indexed[k].0] = avg_rank;
}
i = j;
}
ranks
}
// ── Tests ───────────────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
fn small_config() -> TransformerConfig {
TransformerConfig {
n_subcarriers: 16,
n_keypoints: 17,
d_model: 8,
n_heads: 2,
n_gnn_layers: 1,
}
}
fn small_embed_config() -> EmbeddingConfig {
EmbeddingConfig {
d_model: 8,
d_proj: 128,
temperature: 0.07,
normalize: true,
}
}
fn make_csi(n_pairs: usize, n_sub: usize, seed: u64) -> Vec<Vec<f32>> {
let mut rng = SimpleRng::new(seed);
(0..n_pairs)
.map(|_| (0..n_sub).map(|_| rng.next_f32_unit()).collect())
.collect()
}
// ── ProjectionHead tests ────────────────────────────────────────────
#[test]
fn test_projection_head_output_shape() {
let config = small_embed_config();
let proj = ProjectionHead::new(config);
let input = vec![0.5f32; 8];
let output = proj.forward(&input);
assert_eq!(output.len(), 128);
}
#[test]
fn test_projection_head_l2_normalized() {
let config = small_embed_config();
let proj = ProjectionHead::new(config);
let input = vec![1.0f32; 8];
let output = proj.forward(&input);
let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
assert!(
(norm - 1.0).abs() < 1e-4,
"expected unit norm, got {norm}"
);
}
#[test]
fn test_projection_head_weight_roundtrip() {
let config = small_embed_config();
let proj = ProjectionHead::new(config.clone());
let mut flat = Vec::new();
proj.flatten_into(&mut flat);
assert_eq!(flat.len(), proj.param_count());
let (restored, consumed) = ProjectionHead::unflatten_from(&flat, &config);
assert_eq!(consumed, flat.len());
let input = vec![0.3f32; 8];
let out_orig = proj.forward(&input);
let out_rest = restored.forward(&input);
for (a, b) in out_orig.iter().zip(out_rest.iter()) {
assert!((a - b).abs() < 1e-6, "mismatch: {a} vs {b}");
}
}
// ── InfoNCE loss tests ──────────────────────────────────────────────
#[test]
fn test_info_nce_loss_positive_pairs() {
// Identical embeddings should give low loss (close to log(1) = 0)
let emb = vec![vec![1.0, 0.0, 0.0]; 4];
let loss = info_nce_loss(&emb, &emb, 0.07);
// When all embeddings are identical, all similarities are 1.0,
// so loss = log(N) per sample
let expected = (4.0f32).ln();
assert!(
(loss - expected).abs() < 0.1,
"identical embeddings: expected ~{expected}, got {loss}"
);
}
#[test]
fn test_info_nce_loss_random_pairs() {
// Random embeddings should give higher loss than well-aligned ones
let aligned_a = vec![
vec![1.0, 0.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0, 0.0],
];
let aligned_b = vec![
vec![0.9, 0.1, 0.0, 0.0],
vec![0.1, 0.9, 0.0, 0.0],
];
let random_b = vec![
vec![0.0, 0.0, 1.0, 0.0],
vec![0.0, 0.0, 0.0, 1.0],
];
let loss_aligned = info_nce_loss(&aligned_a, &aligned_b, 0.5);
let loss_random = info_nce_loss(&aligned_a, &random_b, 0.5);
assert!(
loss_random > loss_aligned,
"random should have higher loss: {loss_random} vs {loss_aligned}"
);
}
// ── CsiAugmenter tests ──────────────────────────────────────────────
#[test]
fn test_augmenter_produces_different_views() {
let aug = CsiAugmenter::new();
let csi = vec![vec![1.0f32; 16]; 5];
let (view_a, view_b) = aug.augment_pair(&csi, 42);
// Views should differ (different augmentation pipelines)
let mut any_diff = false;
for (a, b) in view_a.iter().zip(view_b.iter()) {
for (&va, &vb) in a.iter().zip(b.iter()) {
if (va - vb).abs() > 1e-6 {
any_diff = true;
break;
}
}
if any_diff { break; }
}
assert!(any_diff, "augmented views should differ");
}
#[test]
fn test_augmenter_preserves_shape() {
let aug = CsiAugmenter::new();
let csi = vec![vec![0.5f32; 20]; 8];
let (view_a, view_b) = aug.augment_pair(&csi, 99);
assert_eq!(view_a.len(), 8);
assert_eq!(view_b.len(), 8);
for frame in &view_a {
assert_eq!(frame.len(), 20);
}
for frame in &view_b {
assert_eq!(frame.len(), 20);
}
}
// ── EmbeddingExtractor tests ────────────────────────────────────────
#[test]
fn test_embedding_extractor_output_shape() {
let ext = EmbeddingExtractor::new(small_config(), small_embed_config());
let csi = make_csi(4, 16, 42);
let emb = ext.extract(&csi);
assert_eq!(emb.len(), 128);
}
#[test]
fn test_embedding_extractor_weight_roundtrip() {
let ext = EmbeddingExtractor::new(small_config(), small_embed_config());
let weights = ext.flatten_weights();
assert_eq!(weights.len(), ext.param_count());
let mut ext2 = EmbeddingExtractor::new(small_config(), small_embed_config());
ext2.unflatten_weights(&weights).expect("unflatten should succeed");
let csi = make_csi(4, 16, 42);
let emb1 = ext.extract(&csi);
let emb2 = ext2.extract(&csi);
for (a, b) in emb1.iter().zip(emb2.iter()) {
assert!((a - b).abs() < 1e-5, "mismatch: {a} vs {b}");
}
}
// ── FingerprintIndex tests ──────────────────────────────────────────
#[test]
fn test_fingerprint_index_insert_search() {
let mut idx = FingerprintIndex::new(IndexType::EnvironmentFingerprint);
// Insert 10 unit vectors along different axes
for i in 0..10 {
let mut emb = vec![0.0f32; 10];
emb[i] = 1.0;
idx.insert(emb, format!("entry_{i}"), i as u64 * 100);
}
assert_eq!(idx.len(), 10);
// Search for vector close to axis 3
let mut query = vec![0.0f32; 10];
query[3] = 1.0;
let results = idx.search(&query, 3);
assert_eq!(results.len(), 3);
assert_eq!(results[0].entry, 3, "nearest should be entry_3");
assert!(results[0].distance < 0.01, "distance should be ~0");
}
#[test]
fn test_fingerprint_index_anomaly_detection() {
let mut idx = FingerprintIndex::new(IndexType::ActivityPattern);
// Insert clustered embeddings
for i in 0..5 {
let emb = vec![1.0 + i as f32 * 0.01; 8];
idx.insert(emb, format!("normal_{i}"), 0);
}
// Normal query (similar to cluster)
let normal = vec![1.0f32; 8];
assert!(!idx.is_anomaly(&normal, 0.1), "normal should not be anomaly");
// Anomalous query (very different)
let anomaly = vec![-1.0f32; 8];
assert!(idx.is_anomaly(&anomaly, 0.5), "distant should be anomaly");
}
#[test]
fn test_fingerprint_index_types() {
let types = [
IndexType::EnvironmentFingerprint,
IndexType::ActivityPattern,
IndexType::TemporalBaseline,
IndexType::PersonTrack,
];
for &it in &types {
let mut idx = FingerprintIndex::new(it);
idx.insert(vec![1.0, 2.0, 3.0], "test".into(), 0);
assert_eq!(idx.len(), 1);
let results = idx.search(&[1.0, 2.0, 3.0], 1);
assert_eq!(results.len(), 1);
assert!(results[0].distance < 0.01);
}
}
// ── PoseEncoder tests ───────────────────────────────────────────────
#[test]
fn test_pose_encoder_output_shape() {
let enc = PoseEncoder::new(128);
let pose_flat = vec![0.5f32; 51]; // 17 * 3
let out = enc.forward(&pose_flat);
assert_eq!(out.len(), 128);
}
#[test]
fn test_pose_encoder_l2_normalized() {
let enc = PoseEncoder::new(128);
let pose_flat = vec![1.0f32; 51];
let out = enc.forward(&pose_flat);
let norm: f32 = out.iter().map(|x| x * x).sum::<f32>().sqrt();
assert!(
(norm - 1.0).abs() < 1e-4,
"expected unit norm, got {norm}"
);
}
#[test]
fn test_cross_modal_loss_aligned_pairs() {
// Create CSI and pose embeddings that are aligned
let csi_emb = vec![
vec![1.0, 0.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0, 0.0],
vec![0.0, 0.0, 1.0, 0.0],
];
let pose_emb_aligned = vec![
vec![0.95, 0.05, 0.0, 0.0],
vec![0.05, 0.95, 0.0, 0.0],
vec![0.0, 0.05, 0.95, 0.0],
];
let pose_emb_shuffled = vec![
vec![0.0, 0.05, 0.95, 0.0],
vec![0.95, 0.05, 0.0, 0.0],
vec![0.05, 0.95, 0.0, 0.0],
];
let loss_aligned = cross_modal_loss(&csi_emb, &pose_emb_aligned, 0.5);
let loss_shuffled = cross_modal_loss(&csi_emb, &pose_emb_shuffled, 0.5);
assert!(
loss_aligned < loss_shuffled,
"aligned should have lower loss: {loss_aligned} vs {loss_shuffled}"
);
}
// ── Quantized embedding validation ──────────────────────────────────
#[test]
fn test_quantized_embedding_rank_correlation() {
let mut rng = SimpleRng::new(12345);
let embeddings: Vec<Vec<f32>> = (0..20)
.map(|_| (0..32).map(|_| rng.next_gaussian()).collect())
.collect();
let query: Vec<f32> = (0..32).map(|_| rng.next_gaussian()).collect();
let corr = validate_quantized_embeddings(&embeddings, &query, &Quantizer);
assert!(
corr > 0.90,
"rank correlation should be > 0.90, got {corr}"
);
}
// ── Transformer embed() test ────────────────────────────────────────
#[test]
fn test_transformer_embed_shape() {
let t = CsiToPoseTransformer::new(small_config());
let csi = make_csi(4, 16, 42);
let body_feats = t.embed(&csi);
assert_eq!(body_feats.len(), 17);
for f in &body_feats {
assert_eq!(f.len(), 8); // d_model = 8
}
}
}

10
rust-port/wifi-densepose-rs/crates/wifi-densepose-sensing-server/src/graph_transformer.rs
View File

@@ -486,6 +486,16 @@ impl CsiToPoseTransformer {
}
pub fn config(&self) -> &TransformerConfig { &self.config }
/// Extract body-part feature embeddings without regression heads.
/// Returns 17 vectors of dimension d_model (same as forward() but stops
/// before xyz_head/conf_head).
pub fn embed(&self, csi_features: &[Vec<f32>]) -> Vec<Vec<f32>> {
let embedded: Vec<Vec<f32>> = csi_features.iter()
.map(|f| self.csi_embed.forward(f)).collect();
let attended = self.cross_attn.forward(&self.keypoint_queries, &embedded, &embedded);
self.gnn.forward(&attended)
}
/// Collect all trainable parameters into a flat vec.
///
/// Layout: csi_embed | keypoint_queries (flat) | cross_attn | gnn | xyz_head | conf_head

1
rust-port/wifi-densepose-rs/crates/wifi-densepose-sensing-server/src/lib.rs
View File

@@ -12,3 +12,4 @@ pub mod trainer;
pub mod dataset;
pub mod sona;
pub mod sparse_inference;
pub mod embedding;

233
rust-port/wifi-densepose-rs/crates/wifi-densepose-sensing-server/src/main.rs
View File

@@ -13,7 +13,7 @@ mod rvf_pipeline;
mod vital_signs;
// Training pipeline modules (exposed via lib.rs)
use wifi_densepose_sensing_server::{graph_transformer, trainer, dataset};
use wifi_densepose_sensing_server::{graph_transformer, trainer, dataset, embedding};
use std::collections::VecDeque;
use std::net::SocketAddr;
@@ -122,6 +122,22 @@ struct Args {
/// Directory for training checkpoints
#[arg(long, value_name = "DIR")]
checkpoint_dir: Option<PathBuf>,
/// Run self-supervised contrastive pretraining (ADR-024)
#[arg(long)]
pretrain: bool,
/// Number of pretraining epochs (default 50)
#[arg(long, default_value = "50")]
pretrain_epochs: usize,
/// Extract embeddings mode: load model and extract CSI embeddings
#[arg(long)]
embed: bool,
/// Build fingerprint index from embeddings (env|activity|temporal|person)
#[arg(long, value_name = "TYPE")]
build_index: Option<String>,
}
// ── Data types ───────────────────────────────────────────────────────────────
@@ -1536,6 +1552,221 @@ async fn main() {
return;
}
// Handle --pretrain mode: self-supervised contrastive pretraining (ADR-024)
if args.pretrain {
eprintln!("=== WiFi-DensePose Contrastive Pretraining (ADR-024) ===");
let ds_path = args.dataset.clone().unwrap_or_else(|| PathBuf::from("data"));
let source = match args.dataset_type.as_str() {
"wipose" => dataset::DataSource::WiPose(ds_path.clone()),
_ => dataset::DataSource::MmFi(ds_path.clone()),
};
let pipeline = dataset::DataPipeline::new(dataset::DataConfig {
source, ..Default::default()
});
// Generate synthetic or load real CSI windows
let generate_synthetic_windows = || -> Vec<Vec<Vec<f32>>> {
(0..50).map(|i| {
(0..4).map(|a| {
(0..56).map(|s| ((i * 7 + a * 13 + s) as f32 * 0.31).sin() * 0.5).collect()
}).collect()
}).collect()
};
let csi_windows: Vec<Vec<Vec<f32>>> = match pipeline.load() {
Ok(s) if !s.is_empty() => {
eprintln!("Loaded {} samples from {}", s.len(), ds_path.display());
s.into_iter().map(|s| s.csi_window).collect()
}
_ => {
eprintln!("Using synthetic data for pretraining.");
generate_synthetic_windows()
}
};
let n_subcarriers = csi_windows.first()
.and_then(|w| w.first())
.map(|f| f.len())
.unwrap_or(56);
let tf_config = graph_transformer::TransformerConfig {
n_subcarriers, n_keypoints: 17, d_model: 64, n_heads: 4, n_gnn_layers: 2,
};
let transformer = graph_transformer::CsiToPoseTransformer::new(tf_config);
eprintln!("Transformer params: {}", transformer.param_count());
let trainer_config = trainer::TrainerConfig {
epochs: args.pretrain_epochs,
batch_size: 8, lr: 0.001, warmup_epochs: 2, min_lr: 1e-6,
early_stop_patience: args.pretrain_epochs + 1,
pretrain_temperature: 0.07,
..Default::default()
};
let mut t = trainer::Trainer::with_transformer(trainer_config, transformer);
let e_config = embedding::EmbeddingConfig {
d_model: 64, d_proj: 128, temperature: 0.07, normalize: true,
};
let mut projection = embedding::ProjectionHead::new(e_config.clone());
let augmenter = embedding::CsiAugmenter::new();
eprintln!("Starting contrastive pretraining for {} epochs...", args.pretrain_epochs);
let start = std::time::Instant::now();
for epoch in 0..args.pretrain_epochs {
let loss = t.pretrain_epoch(&csi_windows, &augmenter, &mut projection, 0.07, epoch);
if epoch % 10 == 0 || epoch == args.pretrain_epochs - 1 {
eprintln!(" Epoch {epoch}: contrastive loss = {loss:.4}");
}
}
let elapsed = start.elapsed().as_secs_f64();
eprintln!("Pretraining complete in {elapsed:.1}s");
// Save pretrained model as RVF with embedding segment
if let Some(ref save_path) = args.save_rvf {
eprintln!("Saving pretrained model to RVF: {}", save_path.display());
t.sync_transformer_weights();
let weights = t.params().to_vec();
let mut proj_weights = Vec::new();
projection.flatten_into(&mut proj_weights);
let mut builder = RvfBuilder::new();
builder.add_manifest(
"wifi-densepose-pretrained",
env!("CARGO_PKG_VERSION"),
"WiFi DensePose contrastive pretrained model (ADR-024)",
);
builder.add_weights(&weights);
builder.add_embedding(
&serde_json::json!({
"d_model": e_config.d_model,
"d_proj": e_config.d_proj,
"temperature": e_config.temperature,
"normalize": e_config.normalize,
"pretrain_epochs": args.pretrain_epochs,
}),
&proj_weights,
);
match builder.write_to_file(save_path) {
Ok(()) => eprintln!("RVF saved ({} transformer + {} projection params)",
weights.len(), proj_weights.len()),
Err(e) => eprintln!("Failed to save RVF: {e}"),
}
}
return;
}
// Handle --embed mode: extract embeddings from CSI data
if args.embed {
eprintln!("=== WiFi-DensePose Embedding Extraction (ADR-024) ===");
let model_path = match &args.model {
Some(p) => p.clone(),
None => {
eprintln!("Error: --embed requires --model <path> to a pretrained .rvf file");
std::process::exit(1);
}
};
let reader = match RvfReader::from_file(&model_path) {
Ok(r) => r,
Err(e) => { eprintln!("Failed to load model: {e}"); std::process::exit(1); }
};
let weights = reader.weights().unwrap_or_default();
let (embed_config_json, proj_weights) = reader.embedding().unwrap_or_else(|| {
eprintln!("Warning: no embedding segment in RVF, using defaults");
(serde_json::json!({"d_model":64,"d_proj":128,"temperature":0.07,"normalize":true}), Vec::new())
});
let d_model = embed_config_json["d_model"].as_u64().unwrap_or(64) as usize;
let d_proj = embed_config_json["d_proj"].as_u64().unwrap_or(128) as usize;
let tf_config = graph_transformer::TransformerConfig {
n_subcarriers: 56, n_keypoints: 17, d_model, n_heads: 4, n_gnn_layers: 2,
};
let e_config = embedding::EmbeddingConfig {
d_model, d_proj, temperature: 0.07, normalize: true,
};
let mut extractor = embedding::EmbeddingExtractor::new(tf_config, e_config.clone());
// Load transformer weights
if !weights.is_empty() {
if let Err(e) = extractor.transformer.unflatten_weights(&weights) {
eprintln!("Warning: failed to load transformer weights: {e}");
}
}
// Load projection weights
if !proj_weights.is_empty() {
let (proj, _) = embedding::ProjectionHead::unflatten_from(&proj_weights, &e_config);
extractor.projection = proj;
}
// Load dataset and extract embeddings
let _ds_path = args.dataset.clone().unwrap_or_else(|| PathBuf::from("data"));
let csi_windows: Vec<Vec<Vec<f32>>> = (0..10).map(|i| {
(0..4).map(|a| {
(0..56).map(|s| ((i * 7 + a * 13 + s) as f32 * 0.31).sin() * 0.5).collect()
}).collect()
}).collect();
eprintln!("Extracting embeddings from {} CSI windows...", csi_windows.len());
let embeddings = extractor.extract_batch(&csi_windows);
for (i, emb) in embeddings.iter().enumerate() {
let norm: f32 = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
eprintln!(" Window {i}: {d_proj}-dim embedding, ||e|| = {norm:.4}");
}
eprintln!("Extracted {} embeddings of dimension {d_proj}", embeddings.len());
return;
}
// Handle --build-index mode: build a fingerprint index from embeddings
if let Some(ref index_type_str) = args.build_index {
eprintln!("=== WiFi-DensePose Fingerprint Index Builder (ADR-024) ===");
let index_type = match index_type_str.as_str() {
"env" | "environment" => embedding::IndexType::EnvironmentFingerprint,
"activity" => embedding::IndexType::ActivityPattern,
"temporal" => embedding::IndexType::TemporalBaseline,
"person" => embedding::IndexType::PersonTrack,
_ => {
eprintln!("Unknown index type '{}'. Use: env, activity, temporal, person", index_type_str);
std::process::exit(1);
}
};
let tf_config = graph_transformer::TransformerConfig::default();
let e_config = embedding::EmbeddingConfig::default();
let extractor = embedding::EmbeddingExtractor::new(tf_config, e_config);
// Generate synthetic CSI windows for demo
let csi_windows: Vec<Vec<Vec<f32>>> = (0..20).map(|i| {
(0..4).map(|a| {
(0..56).map(|s| ((i * 7 + a * 13 + s) as f32 * 0.31).sin() * 0.5).collect()
}).collect()
}).collect();
let mut index = embedding::FingerprintIndex::new(index_type);
for (i, window) in csi_windows.iter().enumerate() {
let emb = extractor.extract(window);
index.insert(emb, format!("window_{i}"), i as u64 * 100);
}
eprintln!("Built {:?} index with {} entries", index_type, index.len());
// Test a query
let query_emb = extractor.extract(&csi_windows[0]);
let results = index.search(&query_emb, 5);
eprintln!("Top-5 nearest to window_0:");
for r in &results {
eprintln!(" entry={}, distance={:.4}, metadata={}", r.entry, r.distance, r.metadata);
}
return;
}
// Handle --train mode: train a model and exit
if args.train {
eprintln!("=== WiFi-DensePose Training Mode ===");

67
rust-port/wifi-densepose-rs/crates/wifi-densepose-sensing-server/src/rvf_container.rs
View File

@@ -37,6 +37,8 @@ const SEG_META: u8 = 0x07;
const SEG_WITNESS: u8 = 0x0A;
/// Domain profile declarations.
const SEG_PROFILE: u8 = 0x0B;
/// Contrastive embedding model weights and configuration (ADR-024).
pub const SEG_EMBED: u8 = 0x0C;
// ── Pure-Rust CRC32 (IEEE 802.3 polynomial) ────────────────────────────────
@@ -304,6 +306,20 @@ impl RvfBuilder {
self.push_segment(seg_type, payload);
}
/// Add contrastive embedding config and projection head weights (ADR-024).
/// Serializes embedding config as JSON followed by projection weights as f32 LE.
pub fn add_embedding(&mut self, config_json: &serde_json::Value, proj_weights: &[f32]) {
let config_bytes = serde_json::to_vec(config_json).unwrap_or_default();
let config_len = config_bytes.len() as u32;
let mut payload = Vec::with_capacity(4 + config_bytes.len() + proj_weights.len() * 4);
payload.extend_from_slice(&config_len.to_le_bytes());
payload.extend_from_slice(&config_bytes);
for &w in proj_weights {
payload.extend_from_slice(&w.to_le_bytes());
}
self.push_segment(SEG_EMBED, &payload);
}
/// Add witness/proof data as a Witness segment.
pub fn add_witness(&mut self, training_hash: &str, metrics: &serde_json::Value) {
let witness = serde_json::json!({
@@ -528,6 +544,28 @@ impl RvfReader {
.and_then(|data| serde_json::from_slice(data).ok())
}
/// Parse and return the embedding config JSON and projection weights, if present.
pub fn embedding(&self) -> Option<(serde_json::Value, Vec<f32>)> {
let data = self.find_segment(SEG_EMBED)?;
if data.len() < 4 {
return None;
}
let config_len = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
if 4 + config_len > data.len() {
return None;
}
let config: serde_json::Value = serde_json::from_slice(&data[4..4 + config_len]).ok()?;
let weight_data = &data[4 + config_len..];
if weight_data.len() % 4 != 0 {
return None;
}
let weights: Vec<f32> = weight_data
.chunks_exact(4)
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
.collect();
Some((config, weights))
}
/// Number of segments in the container.
pub fn segment_count(&self) -> usize {
self.segments.len()
@@ -911,4 +949,33 @@ mod tests {
assert!(!info.has_quant_info);
assert!(!info.has_witness);
}
#[test]
fn test_rvf_embedding_segment_roundtrip() {
let config = serde_json::json!({
"d_model": 64,
"d_proj": 128,
"temperature": 0.07,
"normalize": true,
});
let weights: Vec<f32> = (0..256).map(|i| (i as f32 * 0.13).sin()).collect();
let mut builder = RvfBuilder::new();
builder.add_manifest("embed-test", "1.0", "embedding test");
builder.add_embedding(&config, &weights);
let data = builder.build();
let reader = RvfReader::from_bytes(&data).unwrap();
assert_eq!(reader.segment_count(), 2);
let (decoded_config, decoded_weights) = reader.embedding()
.expect("embedding segment should be present");
assert_eq!(decoded_config["d_model"], 64);
assert_eq!(decoded_config["d_proj"], 128);
assert!((decoded_config["temperature"].as_f64().unwrap() - 0.07).abs() < 1e-4);
assert_eq!(decoded_weights.len(), weights.len());
for (a, b) in decoded_weights.iter().zip(weights.iter()) {
assert_eq!(a.to_bits(), b.to_bits(), "weight mismatch");
}
}
}

209
rust-port/wifi-densepose-rs/crates/wifi-densepose-sensing-server/src/trainer.rs
View File

@@ -6,6 +6,7 @@
use std::path::Path;
use crate::graph_transformer::{CsiToPoseTransformer, TransformerConfig};
use crate::embedding::{CsiAugmenter, ProjectionHead, info_nce_loss};
use crate::dataset;
/// Standard COCO keypoint sigmas for OKS (17 keypoints).
@@ -18,7 +19,7 @@ pub const COCO_KEYPOINT_SIGMAS: [f32; 17] = [
const SYMMETRY_PAIRS: [(usize, usize); 5] =
[(5, 6), (7, 8), (9, 10), (11, 12), (13, 14)];
/// Individual loss terms from the 6-component composite loss.
/// Individual loss terms from the composite loss (6 supervised + 1 contrastive).
#[derive(Debug, Clone, Default)]
pub struct LossComponents {
pub keypoint: f32,
@@ -27,6 +28,8 @@ pub struct LossComponents {
pub temporal: f32,
pub edge: f32,
pub symmetry: f32,
/// Contrastive loss (InfoNCE); only active during pretraining or when configured.
pub contrastive: f32,
}
/// Per-term weights for the composite loss function.
@@ -38,11 +41,16 @@ pub struct LossWeights {
pub temporal: f32,
pub edge: f32,
pub symmetry: f32,
/// Contrastive loss weight (default 0.0; set >0 for joint training).
pub contrastive: f32,
}
impl Default for LossWeights {
fn default() -> Self {
Self { keypoint: 1.0, body_part: 0.5, uv: 0.5, temporal: 0.1, edge: 0.2, symmetry: 0.1 }
Self {
keypoint: 1.0, body_part: 0.5, uv: 0.5, temporal: 0.1,
edge: 0.2, symmetry: 0.1, contrastive: 0.0,
}
}
}
@@ -124,6 +132,7 @@ pub fn symmetry_loss(kp: &[(f32, f32, f32)]) -> f32 {
pub fn composite_loss(c: &LossComponents, w: &LossWeights) -> f32 {
w.keypoint * c.keypoint + w.body_part * c.body_part + w.uv * c.uv
+ w.temporal * c.temporal + w.edge * c.edge + w.symmetry * c.symmetry
+ w.contrastive * c.contrastive
}
// ── Optimizer ──────────────────────────────────────────────────────────────
@@ -374,6 +383,10 @@ pub struct TrainerConfig {
pub early_stop_patience: usize,
pub checkpoint_every: usize,
pub loss_weights: LossWeights,
/// Contrastive loss weight for joint supervised+contrastive training (default 0.0).
pub contrastive_loss_weight: f32,
/// Temperature for InfoNCE loss during pretraining (default 0.07).
pub pretrain_temperature: f32,
}
impl Default for TrainerConfig {
@@ -382,6 +395,8 @@ impl Default for TrainerConfig {
epochs: 100, batch_size: 32, lr: 0.01, momentum: 0.9, weight_decay: 1e-4,
warmup_epochs: 5, min_lr: 1e-6, early_stop_patience: 10, checkpoint_every: 10,
loss_weights: LossWeights::default(),
contrastive_loss_weight: 0.0,
pretrain_temperature: 0.07,
}
}
}
@@ -546,6 +561,131 @@ impl Trainer {
}
}
/// Run one self-supervised pretraining epoch using SimCLR objective.
/// Does NOT require pose labels -- only CSI windows.
///
/// For each mini-batch:
/// 1. Generate augmented pair (view_a, view_b) for each window
/// 2. Forward each view through transformer to get body_part_features
/// 3. Mean-pool to get frame embedding
/// 4. Project through ProjectionHead
/// 5. Compute InfoNCE loss
/// 6. Estimate gradients via central differences and SGD update
///
/// Returns mean epoch loss.
pub fn pretrain_epoch(
&mut self,
csi_windows: &[Vec<Vec<f32>>],
augmenter: &CsiAugmenter,
projection: &mut ProjectionHead,
temperature: f32,
epoch: usize,
) -> f32 {
if csi_windows.is_empty() {
return 0.0;
}
let lr = self.scheduler.get_lr(epoch);
self.optimizer.set_lr(lr);
let bs = self.config.batch_size.max(1);
let nb = (csi_windows.len() + bs - 1) / bs;
let mut total_loss = 0.0f32;
let tc = self.transformer_config.clone();
let tc_ref = match &tc {
Some(c) => c,
None => return 0.0, // pretraining requires a transformer
};
for bi in 0..nb {
let start = bi * bs;
let end = (start + bs).min(csi_windows.len());
let batch = &csi_windows[start..end];
// Generate augmented pairs and compute embeddings + loss
let snap = self.params.clone();
let mut proj_flat = Vec::new();
projection.flatten_into(&mut proj_flat);
// Combined params: transformer + projection head
let mut combined = snap.clone();
combined.extend_from_slice(&proj_flat);
let t_param_count = snap.len();
let p_config = projection.config.clone();
let tc_c = tc_ref.clone();
let temp = temperature;
// Build augmented views for the batch
let seed_base = (epoch * 10000 + bi) as u64;
let aug_pairs: Vec<_> = batch.iter().enumerate()
.map(|(k, w)| augmenter.augment_pair(w, seed_base + k as u64))
.collect();
// Loss function over combined (transformer + projection) params
let batch_owned: Vec<Vec<Vec<f32>>> = batch.to_vec();
let loss_fn = |params: &[f32]| -> f32 {
let t_params = &params[..t_param_count];
let p_params = &params[t_param_count..];
let mut t = CsiToPoseTransformer::zeros(tc_c.clone());
if t.unflatten_weights(t_params).is_err() {
return f32::MAX;
}
let (proj, _) = ProjectionHead::unflatten_from(p_params, &p_config);
let d = p_config.d_model;
let mut embs_a = Vec::with_capacity(batch_owned.len());
let mut embs_b = Vec::with_capacity(batch_owned.len());
for (k, _w) in batch_owned.iter().enumerate() {
let (ref va, ref vb) = aug_pairs[k];
// Mean-pool body features for view A
let feats_a = t.embed(va);
let mut pooled_a = vec![0.0f32; d];
for f in &feats_a {
for (p, &v) in pooled_a.iter_mut().zip(f.iter()) { *p += v; }
}
let n = feats_a.len() as f32;
if n > 0.0 { for p in pooled_a.iter_mut() { *p /= n; } }
embs_a.push(proj.forward(&pooled_a));
// Mean-pool body features for view B
let feats_b = t.embed(vb);
let mut pooled_b = vec![0.0f32; d];
for f in &feats_b {
for (p, &v) in pooled_b.iter_mut().zip(f.iter()) { *p += v; }
}
let n = feats_b.len() as f32;
if n > 0.0 { for p in pooled_b.iter_mut() { *p /= n; } }
embs_b.push(proj.forward(&pooled_b));
}
info_nce_loss(&embs_a, &embs_b, temp)
};
let batch_loss = loss_fn(&combined);
total_loss += batch_loss;
// Estimate gradient via central differences on combined params
let mut grad = estimate_gradient(&loss_fn, &combined, 1e-4);
clip_gradients(&mut grad, 1.0);
// Update transformer params
self.optimizer.step(&mut self.params, &grad[..t_param_count]);
// Update projection head params
let mut proj_params = proj_flat.clone();
// Simple SGD for projection head
for i in 0..proj_params.len().min(grad.len() - t_param_count) {
proj_params[i] -= lr * grad[t_param_count + i];
}
let (new_proj, _) = ProjectionHead::unflatten_from(&proj_params, &projection.config);
*projection = new_proj;
}
total_loss / nb as f32
}
pub fn checkpoint(&self) -> Checkpoint {
let m = self.history.last().map(|s| s.to_serializable()).unwrap_or(
EpochStatsSerializable {
@@ -713,11 +853,11 @@ mod tests {
assert!(graph_edge_loss(&kp, &[(0,1),(1,2)], &[5.0, 5.0]) < 1e-6);
}
#[test] fn composite_loss_respects_weights() {
let c = LossComponents { keypoint:1.0, body_part:1.0, uv:1.0, temporal:1.0, edge:1.0, symmetry:1.0 };
let w1 = LossWeights { keypoint:1.0, body_part:0.0, uv:0.0, temporal:0.0, edge:0.0, symmetry:0.0 };
let w2 = LossWeights { keypoint:2.0, body_part:0.0, uv:0.0, temporal:0.0, edge:0.0, symmetry:0.0 };
let c = LossComponents { keypoint:1.0, body_part:1.0, uv:1.0, temporal:1.0, edge:1.0, symmetry:1.0, contrastive:0.0 };
let w1 = LossWeights { keypoint:1.0, body_part:0.0, uv:0.0, temporal:0.0, edge:0.0, symmetry:0.0, contrastive:0.0 };
let w2 = LossWeights { keypoint:2.0, body_part:0.0, uv:0.0, temporal:0.0, edge:0.0, symmetry:0.0, contrastive:0.0 };
assert!((composite_loss(&c, &w2) - 2.0 * composite_loss(&c, &w1)).abs() < 1e-6);
let wz = LossWeights { keypoint:0.0, body_part:0.0, uv:0.0, temporal:0.0, edge:0.0, symmetry:0.0 };
let wz = LossWeights { keypoint:0.0, body_part:0.0, uv:0.0, temporal:0.0, edge:0.0, symmetry:0.0, contrastive:0.0 };
assert_eq!(composite_loss(&c, &wz), 0.0);
}
#[test] fn cosine_scheduler_starts_at_initial() {
@@ -878,4 +1018,61 @@ mod tests {
}
}
}
#[test]
fn test_pretrain_epoch_loss_decreases() {
use crate::graph_transformer::{CsiToPoseTransformer, TransformerConfig};
use crate::embedding::{CsiAugmenter, ProjectionHead, EmbeddingConfig};
let tf_config = TransformerConfig {
n_subcarriers: 8, n_keypoints: 17, d_model: 8, n_heads: 2, n_gnn_layers: 1,
};
let transformer = CsiToPoseTransformer::new(tf_config);
let config = TrainerConfig {
epochs: 10, batch_size: 4, lr: 0.001,
warmup_epochs: 0, early_stop_patience: 100,
pretrain_temperature: 0.5,
..Default::default()
};
let mut trainer = Trainer::with_transformer(config, transformer);
let e_config = EmbeddingConfig {
d_model: 8, d_proj: 16, temperature: 0.5, normalize: true,
};
let mut projection = ProjectionHead::new(e_config);
let augmenter = CsiAugmenter::new();
// Synthetic CSI windows (8 windows, each 4 frames of 8 subcarriers)
let csi_windows: Vec<Vec<Vec<f32>>> = (0..8).map(|i| {
(0..4).map(|a| {
(0..8).map(|s| ((i * 7 + a * 3 + s) as f32 * 0.41).sin() * 0.5).collect()
}).collect()
}).collect();
let loss_0 = trainer.pretrain_epoch(&csi_windows, &augmenter, &mut projection, 0.5, 0);
let loss_1 = trainer.pretrain_epoch(&csi_windows, &augmenter, &mut projection, 0.5, 1);
let loss_2 = trainer.pretrain_epoch(&csi_windows, &augmenter, &mut projection, 0.5, 2);
assert!(loss_0.is_finite(), "epoch 0 loss should be finite: {loss_0}");
assert!(loss_1.is_finite(), "epoch 1 loss should be finite: {loss_1}");
assert!(loss_2.is_finite(), "epoch 2 loss should be finite: {loss_2}");
// Loss should generally decrease (or at least the final loss should be less than initial)
assert!(
loss_2 <= loss_0 + 0.5,
"loss should not increase drastically: epoch0={loss_0}, epoch2={loss_2}"
);
}
#[test]
fn test_contrastive_loss_weight_in_composite() {
let c = LossComponents {
keypoint: 0.0, body_part: 0.0, uv: 0.0,
temporal: 0.0, edge: 0.0, symmetry: 0.0, contrastive: 1.0,
};
let w = LossWeights {
keypoint: 0.0, body_part: 0.0, uv: 0.0,
temporal: 0.0, edge: 0.0, symmetry: 0.0, contrastive: 0.5,
};
assert!((composite_loss(&c, &w) - 0.5).abs() < 1e-6);
}
}