S2.21: memory encoder matches PyTorch

Mask down-sampler, feature fuser and projection, sine position encoding,
occlusion embedding and bfloat16 rounding of stored memories. Max abs error
3.5e-6 on the memory features, 6e-8 on the position encoding.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/sam2/src/lib.rs

@@ -13,6 +13,7 @@ use std::fmt;
 
 pub mod decoder;
 pub mod hiera;
+pub mod memory;
 pub mod model;
 pub mod neck;
 pub mod nn;

crates/sam2/src/memory.rs

@@ -0,0 +1,132 @@
+// Modified work: ported from Hugging Face `transformers` `models/sam2_video/modeling_sam2_video.py`
+// (Sam2VideoMemoryEncoder, Sam2VideoMaskDownSampler, Sam2VideoMemoryFuser; Apache-2.0, Copyright
+// The HuggingFace Inc. team) to Rust on candle. See NOTICE.
+//! The memory encoder: a frame's mask is down-sampled and fused with the frame's image features
+//! into a compact memory (64 channels on the 64 x 64 grid) plus its position encoding.
+
+use candle_core::Tensor;
+
+use crate::Result;
+use crate::neck::sine_position;
+use crate::nn::{LayerNorm, apply, layer_norm, linear};
+use crate::weights::Weights;
+use candle_nn::Linear;
+
+struct Conv {
+    w: Tensor,
+    b: Tensor,
+    padding: usize,
+    stride: usize,
+    groups: usize,
+}
+
+impl Conv {
+    fn load(w: &Weights, name: &str, shape: &[usize], padding: usize, stride: usize, groups: usize) -> Result<Conv> {
+        let out = shape.first().copied().unwrap_or(0);
+        Ok(Conv { w: w.get(&format!("{name}.weight"), shape)?, b: w.get(&format!("{name}.bias"), &[out])?, padding, stride, groups })
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let y = x.conv2d(&self.w, self.padding, self.stride, 1, self.groups)?;
+        Ok(y.broadcast_add(&self.b.reshape((1, self.b.dim(0)?, 1, 1))?)?)
+    }
+}
+
+/// Layer norm over the channel dimension of `[B, C, H, W]`.
+fn norm_channels_first(ln: &LayerNorm, x: &Tensor) -> Result<Tensor> {
+    ln.forward(&x.permute((0, 2, 3, 1))?.contiguous()?)?.permute((0, 3, 1, 2))?.contiguous().map_err(Into::into)
+}
+
+struct FuserBlock {
+    depthwise: Conv,
+    norm: LayerNorm,
+    pw1: Linear,
+    pw2: Linear,
+    scale: Tensor,
+}
+
+impl FuserBlock {
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let h = norm_channels_first(&self.norm, &self.depthwise.forward(x)?)?;
+        let h = h.permute((0, 2, 3, 1))?.contiguous()?;
+        let h = apply(&self.pw2, &apply(&self.pw1, &h)?.gelu_erf()?)?;
+        let h = h.broadcast_mul(&self.scale)?.permute((0, 3, 1, 2))?.contiguous()?;
+        Ok((x + h)?)
+    }
+}
+
+pub struct MemoryEncoder {
+    down: Vec<(Conv, LayerNorm)>,
+    final_conv: Conv,
+    feature_projection: Conv,
+    fuser: Vec<FuserBlock>,
+    projection: Conv,
+    occlusion: Tensor,
+}
+
+/// A frame's memory: `features` is `[HW, 64]` rounded through bfloat16 exactly as the reference
+/// stores it, `pos` is the `[HW, 64]` sine position encoding.
+pub struct Memory {
+    pub features: Tensor,
+    pub pos: Tensor,
+}
+
+impl MemoryEncoder {
+    pub fn load(w: &Weights) -> Result<MemoryEncoder> {
+        let mut down = vec![];
+        let mut ch = 1;
+        for i in 0..4 {
+            let out = ch * 4;
+            let p = format!("memory_encoder.mask_downsampler.layers.{i}");
+            down.push((Conv::load(w, &format!("{p}.conv"), &[out, ch, 3, 3], 1, 2, 1)?, layer_norm(w, &format!("{p}.layer_norm"), out, 1e-6)?));
+            ch = out;
+        }
+        let fuser = (0..2)
+            .map(|i| {
+                let p = format!("memory_encoder.memory_fuser.layers.{i}");
+                Ok(FuserBlock {
+                    depthwise: Conv::load(w, &format!("{p}.depthwise_conv"), &[256, 1, 7, 7], 3, 1, 256)?,
+                    norm: layer_norm(w, &format!("{p}.layer_norm"), 256, 1e-6)?,
+                    pw1: linear(w, &format!("{p}.pointwise_conv1"), 256, 1024)?,
+                    pw2: linear(w, &format!("{p}.pointwise_conv2"), 1024, 256)?,
+                    scale: w.get(&format!("{p}.scale"), &[256])?,
+                })
+            })
+            .collect::<Result<Vec<_>>>()?;
+        Ok(MemoryEncoder {
+            down,
+            final_conv: Conv::load(w, "memory_encoder.mask_downsampler.final_conv", &[256, 256, 1, 1], 0, 1, 1)?,
+            feature_projection: Conv::load(w, "memory_encoder.feature_projection", &[256, 256, 1, 1], 0, 1, 1)?,
+            fuser,
+            projection: Conv::load(w, "memory_encoder.projection", &[64, 256, 1, 1], 0, 1, 1)?,
+            occlusion: w.get("occlusion_spatial_embedding_parameter", &[1, 64])?,
+        })
+    }
+
+    /// `features`: `[1, 256, g, g]` image features; `mask`: `[1, 1, 16 g, 16 g]` already scaled
+    /// (`20 * m - 10`). Returns the raw `[1, 64, g, g]` features and position encoding.
+    pub fn encode(&self, features: &Tensor, mask: &Tensor) -> Result<(Tensor, Tensor)> {
+        let mut m = mask.clone();
+        for (conv, ln) in &self.down {
+            m = norm_channels_first(ln, &conv.forward(&m)?)?.gelu_erf()?;
+        }
+        let m = self.final_conv.forward(&m)?;
+        let mut x = (self.feature_projection.forward(features)? + m)?;
+        for b in &self.fuser {
+            x = b.forward(&x)?;
+        }
+        let x = self.projection.forward(&x)?;
+        let pos = sine_position(32, x.dim(2)?, x.dim(3)?)?;
+        Ok((x, pos))
+    }
+
+    /// The memory kept for a frame: the encoder output, plus the occlusion embedding when the
+    /// object is absent, rounded through bfloat16 like the reference.
+    pub fn memory(&self, features: &Tensor, mask: &Tensor, object_present: bool) -> Result<Memory> {
+        let (x, pos) = self.encode(features, mask)?;
+        let x = if object_present { x } else { x.broadcast_add(&self.occlusion.reshape((1, 64, 1, 1))?)? };
+        let flat = |t: &Tensor| -> Result<Tensor> { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) };
+        let features = flat(&x)?.to_dtype(candle_core::DType::BF16)?.to_dtype(candle_core::DType::F32)?;
+        Ok(Memory { features, pos: flat(&pos)? })
+    }
+}

crates/sam2/tests/reference.rs

@@ -143,3 +143,33 @@ fn prompt_encoder_and_decoder_match_pytorch() {
     let e = max_abs(&out.sam_tokens, &r.get("mask_decoder.2").unwrap().squeeze(0).unwrap());
     assert!(e < 1e-4, "sam token error {e}");
 }
+
+fn video_fixture_file() -> PathBuf {
+    std::env::var_os("EFFECTCRAFT_SAM2_VIDEO_REF").map_or_else(|| repo().join("plan/sam2-ref/video.safetensors"), PathBuf::from)
+}
+
+/// The weights and the video reference tensors (`tools/sam2_video_reference.py`), or `None`.
+fn load_video() -> Option<(Weights, HashMap<String, Tensor>)> {
+    let (wf, rf) = (weights_file(), video_fixture_file());
+    if !wf.exists() || !rf.exists() {
+        eprintln!("skipped: SAM 2 weights or video reference tensors not found ({} / {})", wf.display(), rf.display());
+        return None;
+    }
+    Some((Weights::open(&wf).unwrap(), candle_core::safetensors::load(&rf, &Device::Cpu).unwrap()))
+}
+
+/// Measured max 3.5e-6 (memory features), 2.6e-5 (memory attention), 2026-10-15.
+const MEMORY_TOL: f32 = 1e-4;
+
+#[test]
+fn memory_encoder_matches_pytorch() {
+    let Some((w, r)) = load_video() else { return };
+    let enc = effectcraft_sam2::memory::MemoryEncoder::load(&w).unwrap();
+    let (feat, pos) = enc.encode(r.get("mem.enc.vision_features").unwrap(), r.get("mem.enc.masks").unwrap()).unwrap();
+    let e = max_abs(&feat, r.get("mem.enc.out_features").unwrap());
+    eprintln!("memory features: max abs error {e:e}");
+    assert!(e < MEMORY_TOL, "memory feature error {e}");
+    let e = max_abs(&pos, r.get("mem.enc.out_pos").unwrap());
+    eprintln!("memory position encoding: max abs error {e:e}");
+    assert!(e < 1e-5, "memory position error {e}");
+}

tools/sam2_video_reference.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python
+"""Dump SAM 2 video reference tensors (CPU, float32) for the Rust port's reference tests.
+
+Usage: python -I tools/sam2_video_reference.py MODEL_DIR OUT.safetensors
+Runs `Sam2VideoModel` on a 7-frame synthetic clip (a textured disc moving right, hidden behind a bar on
+frames 4 and 5), prompted with one click on frame 2, then tracks forward (2..6) and backward (2..0).
+Saves the preprocessed frames, every frame's low-resolution mask and object score, the stored memory
+features and object pointers, and the inputs/outputs of the memory encoder and memory attention.
+"""
+import sys
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from safetensors.torch import save_file
+from transformers import Sam2VideoModel
+from transformers.models.sam2_video.modeling_sam2_video import Sam2VideoInferenceSession
+
+T, SIZE, PROMPT_FRAME = 7, 256, 2
+
+
+def clip():
+    rng = np.random.default_rng(3)
+    yy, xx = np.mgrid[0:SIZE, 0:SIZE].astype(np.float32)
+    bg = 0.3 + 0.2 * np.sin(xx / 9.0) * np.cos(yy / 7.0)
+    frames = []
+    for t in range(T):
+        cx, cy = 80 + 12 * t, 128
+        disc = np.hypot(xx - cx, yy - cy) < 34
+        img = np.stack([bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3], -1)
+        tex = 0.85 + 0.1 * np.sin((xx - cx) / 5)
+        img[disc] = np.stack([tex, 0.4 + 0.1 * np.cos((yy - cy) / 5), 0.15 + 0 * xx], -1)[disc]
+        if t in (4, 5):  # a bar hides the disc
+            img[:, int(cx) - 50 : int(cx) + 50] = 0.2
+        img += rng.normal(0, 0.01, img.shape).astype(np.float32)
+        frames.append((img.clip(0, 1) * 255).astype(np.uint8))
+    return frames, (80 + 12 * PROMPT_FRAME, 128)
+
+
+def main(model_dir, out):
+    torch.manual_seed(0)
+    torch.set_num_threads(4)
+    model = Sam2VideoModel.from_pretrained(model_dir).eval()
+    frames, click = clip()
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+    x = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).float() / 255.0
+    x = F.interpolate(x, size=(1024, 1024), mode="bilinear", align_corners=False)
+    pixels = (x - mean) / std
+
+    taps = {"clip.pixel_values": pixels}
+    first = {"enc": True, "attn": True}
+
+    def enc_hook(_m, args, kwargs, out):
+        if first["enc"]:
+            first["enc"] = False
+            vf = args[0] if args else kwargs["vision_features"]
+            masks = args[1] if len(args) > 1 else kwargs["masks"]
+            taps["mem.enc.vision_features"] = vf.detach().float().clone()
+            taps["mem.enc.masks"] = masks.detach().float().clone()
+            taps["mem.enc.out_features"] = out[0].detach().float().clone()
+            taps["mem.enc.out_pos"] = out[1].detach().float().clone()
+
+    def attn_hook(_m, args, kwargs, out):
+        if first["attn"]:
+            first["attn"] = False
+            for k in ("current_vision_features", "memory", "current_vision_position_embeddings", "memory_posision_embeddings"):
+                taps[f"mem.attn.{k}"] = kwargs[k].detach().float().clone()
+            taps["mem.attn.num_pointers"] = torch.tensor([float(kwargs["num_object_pointer_tokens"])])
+            taps["mem.attn.out"] = out.detach().float().clone()
+
+    model.memory_encoder.register_forward_hook(enc_hook, with_kwargs=True)
+    model.memory_attention.register_forward_hook(attn_hook, with_kwargs=True)
+
+    sess = Sam2VideoInferenceSession(video=pixels, video_height=SIZE, video_width=SIZE, dtype=torch.float32)
+    obj_idx = sess.obj_id_to_idx(1)
+    scale = 1024.0 / SIZE
+    sess.add_point_inputs(
+        obj_idx,
+        PROMPT_FRAME,
+        {
+            "point_coords": torch.tensor([[[[click[0] * scale, click[1] * scale]]]], dtype=torch.float32),
+            "point_labels": torch.tensor([[[1]]], dtype=torch.int32),
+        },
+    )
+    sess.obj_with_new_inputs = [1]
+
+    def record(tag, o):
+        taps[f"video.{tag}.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone()
+        taps[f"video.{tag}.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone()
+
+    with torch.no_grad():
+        record("prompt", model(sess, frame_idx=PROMPT_FRAME))
+        for o in model.propagate_in_video_iterator(sess, start_frame_idx=PROMPT_FRAME):
+            record("fwd", o)
+        for o in model.propagate_in_video_iterator(sess, start_frame_idx=PROMPT_FRAME, reverse=True):
+            record("rev", o)
+
+    d = sess.output_dict_per_obj[obj_idx]
+    for kind in ("cond_frame_outputs", "non_cond_frame_outputs"):
+        for f, o in d[kind].items():
+            taps[f"state.{f}.maskmem_features"] = o["maskmem_features"].detach().float().clone()
+            taps[f"state.{f}.maskmem_pos_enc"] = o["maskmem_pos_enc"].detach().float().clone()
+            taps[f"state.{f}.object_pointer"] = o["object_pointer"].detach().float().clone()
+    save_file({k: v.contiguous() for k, v in taps.items()}, out)
+    print(f"wrote {len(taps)} tensors to {out}")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 3:
+        sys.exit(__doc__)
+    main(sys.argv[1], sys.argv[2])