diff --git a/v2/crates/wifi-densepose-train/src/losses.rs b/v2/crates/wifi-densepose-train/src/losses.rs index c5a2d29d..1efdd96d 100644 --- a/v2/crates/wifi-densepose-train/src/losses.rs +++ b/v2/crates/wifi-densepose-train/src/losses.rs @@ -118,7 +118,7 @@ impl WiFiDensePoseLoss { // Normalise by number of visible joints in the batch. let n_visible = visibility.sum(Kind::Float); // Guard against division by zero (entire batch may have no labels). - let safe_n = n_visible.clamp(1.0, f64::MAX); + let safe_n = n_visible.clamp_min(1.0); masked.sum(Kind::Float) / safe_n } @@ -165,7 +165,7 @@ impl WiFiDensePoseLoss { let masked_target_uv = target_uv * &fg_mask_f; // Count foreground pixels × 48 channels to normalise. - let n_fg = fg_mask_f.sum(Kind::Float).clamp(1.0, f64::MAX); + let n_fg = fg_mask_f.sum(Kind::Float).clamp_min(1.0); // Smooth-L1 with beta=1.0, reduction=Sum then divide by fg count. let uv_loss_sum = masked_pred_uv.smooth_l1_loss(&masked_target_uv, Reduction::Sum, 1.0); @@ -234,7 +234,7 @@ impl WiFiDensePoseLoss { // UV loss (foreground masked) let fg_mask = target_int.not_equal(0_i64); let fg_mask_f = fg_mask.unsqueeze(1).expand_as(pu).to_kind(Kind::Float); - let n_fg = fg_mask_f.sum(Kind::Float).clamp(1.0, f64::MAX); + let n_fg = fg_mask_f.sum(Kind::Float).clamp_min(1.0); let uv_loss = (pu * &fg_mask_f).smooth_l1_loss(&(tu * &fg_mask_f), Reduction::Sum, 1.0) / n_fg; @@ -743,10 +743,11 @@ mod tests { } // Visible batch (index 1) should have non-zero heatmaps. + let heatmaps_ref = &heatmaps; let batch1_sum: f32 = (0..num_joints) .map(|j| { (0..size) - .flat_map(|r| (0..size).map(move |c| heatmaps[[1, j, r, c]])) + .flat_map(|r| (0..size).map(move |c| heatmaps_ref[[1, j, r, c]])) .sum::() }) .sum(); diff --git a/v2/crates/wifi-densepose-train/src/metrics.rs b/v2/crates/wifi-densepose-train/src/metrics.rs index 913afa05..e5988d1e 100644 --- a/v2/crates/wifi-densepose-train/src/metrics.rs +++ b/v2/crates/wifi-densepose-train/src/metrics.rs @@ -19,6 +19,7 @@ use ndarray::{Array1, Array2, ArrayView1, ArrayView2}; use petgraph::graph::{DiGraph, NodeIndex}; +use petgraph::visit::EdgeRef; use ruvector_mincut::{DynamicMinCut, MinCutBuilder}; use std::collections::VecDeque; @@ -106,6 +107,24 @@ impl Default for MetricsResult { } } +// --------------------------------------------------------------------------- +// EvalMetrics +// --------------------------------------------------------------------------- + +/// Per-evaluation pose metrics. +/// +/// Plain value container produced by evaluation runs: lower `mpjpe`/`gps` +/// and higher `pck_at_05` indicate better predictions. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct EvalMetrics { + /// Mean Per-Joint Position Error (normalised units). + pub mpjpe: f64, + /// Percentage of Correct Keypoints at threshold 0.05 (0-1 scale). + pub pck_at_05: f64, + /// Geodesic Point Similarity error for DensePose surface predictions. + pub gps: f64, +} + // --------------------------------------------------------------------------- // MetricsAccumulator // --------------------------------------------------------------------------- diff --git a/v2/crates/wifi-densepose-train/src/model.rs b/v2/crates/wifi-densepose-train/src/model.rs index 484eb478..d06ebd01 100644 --- a/v2/crates/wifi-densepose-train/src/model.rs +++ b/v2/crates/wifi-densepose-train/src/model.rs @@ -182,7 +182,7 @@ impl WiFiDensePoseModel { self.vs .trainable_variables() .iter() - .map(|t| t.numel()) + .map(|t| t.numel() as i64) .sum() } @@ -297,7 +297,12 @@ fn apply_antenna_attention(x: &Tensor, lambda: f32) -> Tensor { let xi = x.select(0, bi as i64); // [n_ant, n_sc] // Move to CPU and convert to f32 for the pure-Rust attention kernel. - let flat: Vec = Vec::from(xi.to_kind(Kind::Float).to_device(Device::Cpu).contiguous()); + let flat: Vec = Vec::::try_from( + xi.to_kind(Kind::Float) + .to_device(Device::Cpu) + .flatten(0, -1), + ) + .expect("antenna tensor to vec"); // Q = K = V = the antenna features (self-attention over antenna paths). let out = attn_mincut( @@ -350,7 +355,12 @@ fn apply_spatial_attention(x: &Tensor) -> Tensor { for bi in 0..b { // Extract [C, H*W] and transpose to [H*W, C]. let xi = x.select(0, bi).reshape([c, h * w]).transpose(0, 1); // [H*W, C] - let flat: Vec = Vec::from(xi.to_kind(Kind::Float).to_device(Device::Cpu).contiguous()); + let flat: Vec = Vec::::try_from( + xi.to_kind(Kind::Float) + .to_device(Device::Cpu) + .flatten(0, -1), + ) + .expect("spatial tensor to vec"); // Build token slices — one per spatial position. let tokens: Vec<&[f32]> = (0..n_spatial).map(|i| &flat[i * d..(i + 1) * d]).collect(); diff --git a/v2/crates/wifi-densepose-train/src/proof.rs b/v2/crates/wifi-densepose-train/src/proof.rs index 35f9ff14..b6114e4a 100644 --- a/v2/crates/wifi-densepose-train/src/proof.rs +++ b/v2/crates/wifi-densepose-train/src/proof.rs @@ -153,11 +153,11 @@ pub fn run_proof(proof_dir: &Path) -> Result = Vec::::from(kp.to_kind(Kind::Double).flatten(0, -1)) + let kp_vec: Vec = Vec::::try_from(kp.to_kind(Kind::Double).flatten(0, -1))? .iter() .map(|&x| x as f32) .collect(); - let vis_vec: Vec = Vec::::from(vis.to_kind(Kind::Double).flatten(0, -1)) + let vis_vec: Vec = Vec::::try_from(vis.to_kind(Kind::Double).flatten(0, -1))? .iter() .map(|&x| x as f32) .collect(); @@ -261,7 +261,7 @@ pub fn hash_model_weights(model: &WiFiDensePoseModel) -> String { .flatten(0, -1) .to_kind(Kind::Float) .to_device(Device::Cpu); - let values: Vec = Vec::::from(&flat); + let values: Vec = Vec::::try_from(&flat).expect("param tensor to vec"); let mut buf = vec![0u8; values.len() * 4]; for (i, v) in values.iter().enumerate() { let bytes = v.to_le_bytes(); @@ -292,6 +292,15 @@ pub fn load_expected_hash(proof_dir: &Path) -> Result, std::io::E Ok(if hash.is_empty() { None } else { Some(hash) }) } +/// Verify that `path` is a valid checkpoint directory. +/// +/// Returns `true` only when the path exists and is a directory. Deterministic +/// and side-effect free — repeated calls always return the same result for an +/// unchanged filesystem. +pub fn verify_checkpoint_dir(path: &Path) -> bool { + path.is_dir() +} + /// Save the expected model hash to `/expected_proof.sha256`. /// /// Creates `proof_dir` if it does not already exist. diff --git a/v2/crates/wifi-densepose-train/src/trainer.rs b/v2/crates/wifi-densepose-train/src/trainer.rs index cbb7da72..a022fed0 100644 --- a/v2/crates/wifi-densepose-train/src/trainer.rs +++ b/v2/crates/wifi-densepose-train/src/trainer.rs @@ -582,11 +582,13 @@ fn kp_to_heatmap_tensor( let num_kp = kp_tensor.size()[1] as usize; // Convert to ndarray for generate_target_heatmaps. - let kp_vec: Vec = Vec::::from(kp_tensor.to_kind(Kind::Double).flatten(0, -1)) + let kp_vec: Vec = Vec::::try_from(kp_tensor.to_kind(Kind::Double).flatten(0, -1)) + .expect("kp tensor to vec") .iter() .map(|&x| x as f32) .collect(); - let vis_vec: Vec = Vec::::from(vis_tensor.to_kind(Kind::Double).flatten(0, -1)) + let vis_vec: Vec = Vec::::try_from(vis_tensor.to_kind(Kind::Double).flatten(0, -1)) + .expect("vis tensor to vec") .iter() .map(|&x| x as f32) .collect(); @@ -622,8 +624,8 @@ fn heatmap_to_keypoints(heatmaps: &Tensor) -> Tensor { let arg = flat.argmax(-1, false); // Decompose linear index into (row, col). - let row = (&arg / w).to_kind(Kind::Float); // [B, 17] - let col = (&arg % w).to_kind(Kind::Float); // [B, 17] + let row = arg.divide_scalar_mode(w, "floor").to_kind(Kind::Float); // [B, 17] + let col = arg.remainder(w).to_kind(Kind::Float); // [B, 17] // Normalize to [0, 1] let x = col / (w - 1) as f64; @@ -639,7 +641,8 @@ fn heatmap_to_keypoints(heatmaps: &Tensor) -> Tensor { fn extract_kp_ndarray(kp_tensor: &Tensor, batch_idx: usize) -> Array2 { let num_kp = kp_tensor.size()[1] as usize; let row = kp_tensor.select(0, batch_idx as i64); - let data: Vec = Vec::::from(row.to_kind(Kind::Double).flatten(0, -1)) + let data: Vec = Vec::::try_from(row.to_kind(Kind::Double).flatten(0, -1)) + .expect("kp tensor to vec") .iter() .map(|&v| v as f32) .collect(); @@ -652,7 +655,8 @@ fn extract_kp_ndarray(kp_tensor: &Tensor, batch_idx: usize) -> Array2 { fn extract_vis_ndarray(vis_tensor: &Tensor, batch_idx: usize) -> Array1 { let num_kp = vis_tensor.size()[1] as usize; let row = vis_tensor.select(0, batch_idx as i64); - let data: Vec = Vec::::from(row.to_kind(Kind::Double)) + let data: Vec = Vec::::try_from(row.to_kind(Kind::Double)) + .expect("vis tensor to vec") .iter() .map(|&v| v as f32) .collect(); 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 e20da273..fa408e74 100644 --- a/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs +++ b/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs @@ -283,7 +283,9 @@ mod tests { 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"); + // safetensors, not .pt: this torch build's _save_parameters/_load_parameters + // .pt roundtrip is broken on Windows (GenericDict internal assert) + let path = tmp.path().join("wiflow_std.safetensors"); model.save(&path).expect("save"); model.load(&path).expect("load"); let out = model.forward_inference(&random_csi(&cfg, 1));