From d2291926f018d7eb078e2bc7140a39dd46be35bd Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:16:53 +0200 Subject: [PATCH] S2.12: FPN neck and vision encoder match PyTorch Max abs error 9.7e-6 on the features, 6e-8 on the sine position encodings. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/neck.rs | 141 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 30 +++++++ 3 files changed, 172 insertions(+) create mode 100644 crates/sam2/src/neck.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 246739e45..6515eff7b 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -12,6 +12,7 @@ use std::fmt; pub mod hiera; +pub mod neck; pub mod nn; pub mod weights; diff --git a/crates/sam2/src/neck.rs b/crates/sam2/src/neck.rs new file mode 100644 index 000000000..0a82ec7a9 --- /dev/null +++ b/crates/sam2/src/neck.rs @@ -0,0 +1,141 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py` +// (Sam2VisionNeck, Sam2SinePositionEmbedding, Sam2VisionModel; Apache-2.0, Copyright The +// HuggingFace Inc. team) to Rust on candle. See NOTICE. +//! The feature pyramid neck on top of Hiera, and the vision encoder that joins them. + +use candle_core::{Device, Tensor}; + +use crate::hiera::{Hiera, HieraConfig}; +use crate::weights::Weights; +use crate::{Error, Result}; + +/// Sine position encoding `[1, 2 * num_pos_features, h, w]`, normalised (`Sam2SinePositionEmbedding`). +pub fn sine_position(num_pos_features: usize, h: usize, w: usize) -> Result { + let eps = 1e-6f32; + let scale = 2.0 * std::f32::consts::PI; + let dim_t: Vec = (0..num_pos_features).map(|i| 10000f32.powf(2.0 * (i / 2) as f32 / num_pos_features as f32)).collect(); + let axis = |n: usize, k: usize| -> Vec { + // value of one coordinate at index k of an axis of length n, over the feature dimension + let e = (k + 1) as f32 / (n as f32 + eps) * scale; + dim_t.iter().enumerate().map(|(i, d)| if i % 2 == 0 { (e / d).sin() } else { (e / d).cos() }).collect() + }; + let (ys, xs): (Vec<_>, Vec<_>) = ((0..h).map(|k| axis(h, k)).collect(), (0..w).map(|k| axis(w, k)).collect()); + let mut out = vec![0f32; 2 * num_pos_features * h * w]; + for y in 0..h { + for x in 0..w { + for f in 0..num_pos_features { + let (py, px) = (ys.get(y).and_then(|v| v.get(f)), xs.get(x).and_then(|v| v.get(f))); + if let (Some(py), Some(px)) = (py, px) { + if let Some(o) = out.get_mut((f * h + y) * w + x) { + *o = *py; + } + if let Some(o) = out.get_mut(((num_pos_features + f) * h + y) * w + x) { + *o = *px; + } + } + } + } + } + Ok(Tensor::from_vec(out, (1, 2 * num_pos_features, h, w), &Device::Cpu)?) +} + +struct Conv1x1 { + w: Tensor, + b: Tensor, + out: usize, +} + +impl Conv1x1 { + fn forward(&self, x: &Tensor) -> Result { + Ok(x.conv2d(&self.w, 0, 1, 1, 1)?.broadcast_add(&self.b.reshape((1, self.out, 1, 1))?)?) + } +} + +/// The FPN neck: a 1x1 conv per stage, top-down additions on the two lowest-resolution levels. +pub struct Neck { + convs: Vec, + top_down: Vec, + pos_features: usize, +} + +/// What the neck produces, both in the neck's own order (lowest resolution first, all four +/// levels) as `[B, C, H, W]`. +pub struct NeckOutput { + pub features: Vec, + pub positions: Vec, +} + +impl Neck { + pub fn load(w: &Weights, channels: &[usize], fpn: usize, top_down: &[usize]) -> Result { + let convs = channels + .iter() + .enumerate() + .map(|(j, c)| { + Ok(Conv1x1 { + w: w.get(&format!("vision_encoder.neck.convs.{j}.weight"), &[fpn, *c, 1, 1])?, + b: w.get(&format!("vision_encoder.neck.convs.{j}.bias"), &[fpn])?, + out: fpn, + }) + }) + .collect::>>()?; + Ok(Neck { convs, top_down: top_down.to_vec(), pos_features: fpn / 2 }) + } + + /// `stages`: Hiera stage outputs, NHWC, highest resolution first. + pub fn forward(&self, stages: &[Tensor]) -> Result { + let n = self.convs.len().checked_sub(1).ok_or_else(|| Error::Weights("neck has no convolutions".into()))?; + if stages.len() != self.convs.len() { + return Err(Error::Input(format!("neck expects {} stages, got {}", self.convs.len(), stages.len()))); + } + let (mut features, mut positions) = (vec![], vec![]); + let mut prev: Option = None; + for i in (0..=n).rev() { + let stage = stages.get(i).ok_or_else(|| Error::Input("missing stage".into()))?; + let conv = self.convs.get(n - i).ok_or_else(|| Error::Weights("missing neck convolution".into()))?; + let lateral = conv.forward(&stage.permute((0, 3, 1, 2))?.contiguous()?)?; + let cur = match (&prev, self.top_down.contains(&i) && i != n) { + (Some(p), true) => { + let (h, w) = (p.dim(2)? * 2, p.dim(3)? * 2); + (lateral + p.upsample_nearest2d(h, w)?)? + } + _ => lateral, + }; + positions.push(sine_position(self.pos_features, cur.dim(2)?, cur.dim(3)?)?); + features.push(cur.clone()); + prev = Some(cur); + } + Ok(NeckOutput { features, positions }) + } +} + +/// Hiera plus neck: `Sam2VisionModel`. +pub struct VisionEncoder { + pub hiera: Hiera, + pub neck: Neck, + levels: usize, +} + +/// The multi-scale image features the decoder and memory modules use, highest resolution first +/// (256², 128², 64² for a 1024 input), `[B, C, H, W]`. +pub struct ImageFeatures { + pub features: Vec, + pub positions: Vec, +} + +impl VisionEncoder { + pub fn load_tiny(w: &Weights) -> Result { + let cfg = HieraConfig::tiny(); + Ok(VisionEncoder { hiera: Hiera::load(w, &cfg)?, neck: Neck::load(w, &[768, 384, 192, 96], 256, &[2, 3])?, levels: 3 }) + } + + pub fn forward(&self, pixels: &Tensor) -> Result { + let stages = self.hiera.forward(pixels)?; + let out = self.neck.forward(&stages)?; + // keep the last `levels` of the neck order and reverse: high to low resolution + let keep = |v: Vec| -> Vec { + let skip = v.len().saturating_sub(self.levels); + v.into_iter().skip(skip).rev().collect() + }; + Ok(ImageFeatures { features: keep(out.features), positions: keep(out.positions) }) + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index e5b7ae596..bae7d42e3 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -77,3 +77,33 @@ fn hiera_is_deterministic_and_rejects_bad_sizes() { let odd = pixels.narrow(2, 0, 1022).unwrap(); assert!(hiera.forward(&odd).is_err()); } + +/// The neck adds 1x1 convolutions and nearest-neighbour sums to the Hiera output; the sine +/// position encodings are exact functions of the grid size. +const NECK_TOL: f32 = 1e-4; + +#[test] +fn neck_and_vision_encoder_match_pytorch() { + let Some((w, r)) = load() else { return }; + let enc = effectcraft_sam2::neck::VisionEncoder::load_tiny(&w).unwrap(); + let pixels = r.get("input.pixel_values").unwrap(); + let stages = enc.hiera.forward(pixels).unwrap(); + let out = enc.neck.forward(&stages).unwrap(); + assert_eq!((out.features.len(), out.positions.len()), (4, 4)); + for (i, f) in out.features.iter().enumerate() { + let e = max_abs(f, r.get(&format!("vision_encoder.neck.{i}")).unwrap()); + eprintln!("neck feature {i}: max abs error {e:e}"); + assert!(e < NECK_TOL, "neck feature {i} error {e}"); + } + for (i, p) in out.positions.iter().enumerate() { + let e = max_abs(p, r.get(&format!("vision_encoder.neck.{}", i + 4)).unwrap()); + eprintln!("neck position {i}: max abs error {e:e}"); + assert!(e < NECK_TOL, "neck position {i} error {e}"); + } + // The vision encoder keeps the three highest-resolution levels, highest first. + let feats = enc.forward(pixels).unwrap(); + let dims: Vec = feats.features.iter().map(|f| f.dim(2).unwrap()).collect(); + assert_eq!(dims, vec![256, 128, 64]); + let e = max_abs(feats.features.first().unwrap(), r.get("vision_encoder.neck.3").unwrap()); + assert!(e < NECK_TOL); +}