From e4f4fea0f238ee111be2b64163f4cfcc6c6c3dc1 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:30:33 +0200 Subject: [PATCH] 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 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/memory.rs | 132 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 30 ++++++++ tools/sam2_video_reference.py | 112 ++++++++++++++++++++++++++++ 4 files changed, 275 insertions(+) create mode 100644 crates/sam2/src/memory.rs create mode 100644 tools/sam2_video_reference.py diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 0a67f8ad7..18fe5a8ae 100644 --- a/crates/sam2/src/lib.rs +++ b/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; diff --git a/crates/sam2/src/memory.rs b/crates/sam2/src/memory.rs new file mode 100644 index 000000000..51e3d3c92 --- /dev/null +++ b/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 { + 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 { + 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 { + 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 { + 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, + 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 { + 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::>>()?; + 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 { + 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 { 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)? }) + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 6536a9eb5..f90d03ab8 100644 --- a/crates/sam2/tests/reference.rs +++ b/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)> { + 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}"); +} diff --git a/tools/sam2_video_reference.py b/tools/sam2_video_reference.py new file mode 100644 index 000000000..6e769d955 --- /dev/null +++ b/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])