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 <noreply@anthropic.com>

crates/sam2/src/lib.rs

@@ -12,6 +12,7 @@
 use std::fmt;
 
 pub mod hiera;
+pub mod neck;
 pub mod nn;
 pub mod weights;
 

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<Tensor> {
+    let eps = 1e-6f32;
+    let scale = 2.0 * std::f32::consts::PI;
+    let dim_t: Vec<f32> = (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<f32> {
+        // 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<Tensor> {
+        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<Conv1x1>,
+    top_down: Vec<usize>,
+    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<Tensor>,
+    pub positions: Vec<Tensor>,
+}
+
+impl Neck {
+    pub fn load(w: &Weights, channels: &[usize], fpn: usize, top_down: &[usize]) -> Result<Neck> {
+        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::<Result<Vec<_>>>()?;
+        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<NeckOutput> {
+        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<Tensor> = 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<Tensor>,
+    pub positions: Vec<Tensor>,
+}
+
+impl VisionEncoder {
+    pub fn load_tiny(w: &Weights) -> Result<VisionEncoder> {
+        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<ImageFeatures> {
+        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<Tensor>| -> Vec<Tensor> {
+            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) })
+    }
+}

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<usize> = 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);
+}