mirror of
https://github.com/storytold/effectcraft.git
synced 2026-10-31 21:48:20 -06:00
S2.63: GPU memory attention and memory encoder; one shared device; CPU fallback
Memory attention (1.8e-5 vs PyTorch) and the memory encoder (2.6e-6) as wgpu kernels, convolution and layer-norm+GELU kernels, VideoModel::load_auto picks the GPU when there is one (EFFECTCRAFT_SAM2_DEVICE overrides) and falls back to the CPU on any GPU error; disable_gpu switches it off mid-clip. One process-wide device (several at once crashed the Vulkan loader). A tracked frame: 4.3 s CPU, 0.7 s on the RX 7900 GRE. Clips tracked on the GPU match Sam2VideoModel (worst logit error 0.0032). Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Kristian Myram
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
5c7eb3c2ae
commit
ee3498b24d
@@ -478,7 +478,7 @@ fn status(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
.keys()
|
||||
.filter_map(|f| chain.trust_at(*f).map(|t| (f.to_string(), json!({"confidence": t.confidence, "occluded": t.occluded}))))
|
||||
.collect();
|
||||
out["tracker"] = json!({"model": s.models_roto_active(), "saved": chain.saved_model(), "frames": tracked});
|
||||
out["tracker"] = json!({"model": s.models_roto_active(), "device": fx::installed_model().and_then(|m| m.device_name()), "saved": chain.saved_model(), "frames": tracked});
|
||||
let notice = s.roto_model_notice(s.uses_roto_model(cid, lid, uid)).or_else(|| s.saved_tracking_notice(chain.saved_model()));
|
||||
out["notice"] = json!(notice);
|
||||
out["source"] = json!(if tracked.contains_key(&frame.to_string()) {
|
||||
|
||||
@@ -70,6 +70,57 @@ fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let _ = attn.forward(&cur, &cur_pos, &mem, &mem_pos, 64)?;
|
||||
println!("memory attention (run {run}, {m} memory tokens): {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
||||
}
|
||||
// The same two costs on the GPU (wgpu kernels), when there is one.
|
||||
match effectcraft_sam2::gpu::Gpu::shared() {
|
||||
Ok(gpu) => {
|
||||
println!("GPU: {}", gpu.adapter_name());
|
||||
let enc = effectcraft_sam2::gpu::encoder::GpuEncoder::load(gpu.clone(), &weights, model.vision.hiera.clone())?;
|
||||
let gattn = effectcraft_sam2::gpu::memattn::GpuMemoryAttention::load(gpu, &weights)?;
|
||||
for run in 0..3 {
|
||||
let t = Instant::now();
|
||||
let _ = enc.encode(&pixels)?;
|
||||
let encode = t.elapsed().as_secs_f64() * 1e3;
|
||||
let t = Instant::now();
|
||||
let _ = gattn.forward(&cur, &cur_pos, &mem, &mem_pos, 64)?;
|
||||
println!("GPU run {run}: encode {encode:.0} ms, memory attention (full bank) {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
||||
}
|
||||
}
|
||||
Err(e) => println!("no GPU: {e}"),
|
||||
}
|
||||
// A whole tracked frame (encode, memory attention, decode, memory encoding) through the video session, on each backend.
|
||||
struct Same(Tensor);
|
||||
impl effectcraft_sam2::session::FrameSource for Same {
|
||||
fn num_frames(&self) -> usize {
|
||||
12
|
||||
}
|
||||
fn frame(&self, _i: usize) -> effectcraft_sam2::Result<Tensor> {
|
||||
Ok(self.0.clone())
|
||||
}
|
||||
}
|
||||
for backend in [effectcraft_sam2::session::Backend::Cpu, effectcraft_sam2::session::Backend::Gpu] {
|
||||
let Ok(vm) = effectcraft_sam2::session::VideoModel::load_with(&weights, backend) else { continue };
|
||||
let vm = std::sync::Arc::new(vm);
|
||||
let mut tracker = effectcraft_sam2::session::VideoTracker::new(vm.clone(), Box::new(Same(pixels.clone())))?;
|
||||
let obj = effectcraft_sam2::session::ObjectId(1);
|
||||
let t = Instant::now();
|
||||
tracker.prompt(0, obj, &Prompt { points: vec![([512.0, 512.0], true)], bbox: None })?;
|
||||
println!("{:?} ({}): prompt frame {:.0} ms", backend, vm.device_name(), t.elapsed().as_secs_f64() * 1e3);
|
||||
let mut last = Instant::now();
|
||||
let mut times = vec![];
|
||||
tracker.propagate(0, effectcraft_sam2::session::Direction::Forward, None, &effectcraft_sam2::session::Cancel::new(), &mut |m| {
|
||||
if m.frame > 0 {
|
||||
times.push(last.elapsed().as_secs_f64() * 1e3);
|
||||
}
|
||||
last = Instant::now();
|
||||
})?;
|
||||
println!(
|
||||
"{:?}: tracked frames {:.0} ms each (first {:.0}, last {:.0})",
|
||||
backend,
|
||||
times.iter().sum::<f64>() / times.len().max(1) as f64,
|
||||
times.first().copied().unwrap_or(0.0),
|
||||
times.last().copied().unwrap_or(0.0)
|
||||
);
|
||||
}
|
||||
println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get()));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -107,7 +107,7 @@ impl GpuEncoder {
|
||||
|
||||
fn norm(&self, x: &Buf, rows: usize, dim: usize, n: &GNorm) -> Result<Buf> {
|
||||
let out = self.gpu.zeros(rows * dim)?;
|
||||
self.gpu.layernorm(x, &n.w, &n.b, &out, rows, dim, self.eps)?;
|
||||
self.gpu.layernorm(x, &n.w, &n.b, &out, rows, dim, self.eps, false)?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
@@ -227,6 +227,7 @@ impl GpuEncoder {
|
||||
/// Encode normalised `[1, 3, H, W]` pixels (`H` and `W` multiples of 32).
|
||||
pub fn encode(&self, pixels: &Tensor) -> Result<GpuFeatures> {
|
||||
let g = &self.gpu;
|
||||
let _turn = g.exclusive();
|
||||
let (_, c, ph, pw) = pixels.dims4()?;
|
||||
let window0 = self.cfg.windows.first().copied().unwrap_or(8);
|
||||
if c != 3 || ph == 0 || pw == 0 || ph % 4 != 0 || pw % 4 != 0 || (ph / 4) % window0 != 0 || (pw / 4) % window0 != 0 {
|
||||
|
||||
@@ -29,7 +29,7 @@ fn reduce_max(li: u32) -> f32 {
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- layer norm over the last dimension
|
||||
// words: 0 rows, 1 cols, 2 eps (f32). in0 x, in1 weight, in2 bias, out0 y. One workgroup per row.
|
||||
// words: 0 rows, 1 cols, 2 eps (f32), 3 gelu (1: apply GELU to the result). in0 x, in1 weight, in2 bias, out0 y. One workgroup per row.
|
||||
@compute @workgroup_size(64)
|
||||
fn layernorm(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
|
||||
let row = wg.y * nw.x + wg.x;
|
||||
@@ -44,7 +44,40 @@ fn layernorm(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw:
|
||||
for (var c = li.x; c < cols; c += 64u) { let d = in0[base + c] - mean; q += d * d; }
|
||||
red[li.x] = q;
|
||||
let inv = inverseSqrt(reduce_sum(li.x) / f32(cols) + f(2u));
|
||||
for (var c = li.x; c < cols; c += 64u) { out0[base + c] = (in0[base + c] - mean) * inv * in1[c] + in2[c]; }
|
||||
for (var c = li.x; c < cols; c += 64u) {
|
||||
var y = (in0[base + c] - mean) * inv * in1[c] + in2[c];
|
||||
if (u(3u) == 1u) { y = gelu_(y); }
|
||||
out0[base + c] = y;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- convolution over channels-last [H, W, Ci]
|
||||
// words: 0 n (output elements), 1 H, 2 W, 3 Ci, 4 Co, 5 Ho, 6 Wo, 7 k, 8 stride, 9 pad, 10 groups.
|
||||
// in0 input, in1 weight [Co, Ci / groups, k, k], in2 bias [Co]; output [Ho, Wo, Co].
|
||||
@compute @workgroup_size(256)
|
||||
fn conv2d(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
|
||||
let i = gid1(wg, nw, li);
|
||||
if (i >= u(0u)) { return; }
|
||||
let H = u(1u); let W = u(2u); let Ci = u(3u); let Co = u(4u); let Wo = u(6u);
|
||||
let k = u(7u); let stride = u(8u); let pad = i32(u(9u)); let groups = u(10u);
|
||||
let co = i % Co; let t = i / Co;
|
||||
let ox = t % Wo; let oy = t / Wo;
|
||||
let cig = Ci / groups;
|
||||
let g0 = co / (Co / groups);
|
||||
var acc = in2[co];
|
||||
for (var c = 0u; c < cig; c++) {
|
||||
let ci = g0 * cig + c;
|
||||
for (var ky = 0u; ky < k; ky++) {
|
||||
let iy = i32(oy * stride + ky) - pad;
|
||||
if (iy < 0 || iy >= i32(H)) { continue; }
|
||||
for (var kx = 0u; kx < k; kx++) {
|
||||
let ix = i32(ox * stride + kx) - pad;
|
||||
if (ix < 0 || ix >= i32(W)) { continue; }
|
||||
acc += in0[(u32(iy) * W + u32(ix)) * Ci + ci] * in1[((co * cig + c) * k + ky) * k + kx];
|
||||
}
|
||||
}
|
||||
}
|
||||
out0[i] = acc;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- softmax over rows, in place on out0
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
// Original work for EffectCraft; the maths follows the Hugging Face `transformers` SAM 2 video
|
||||
// implementation (Apache-2.0, see NOTICE).
|
||||
//! Memory attention on the GPU: four layers of self-attention over the current frame and cross-attention to
|
||||
//! the memory bank, with axial rotary position encoding. The attention is one head of 256 channels; the
|
||||
//! score matrix against a full memory bank is 4096 x 28736 floats, which stays on the GPU.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use candle_core::{Device, Tensor};
|
||||
|
||||
use super::{Act, Buf, GLinear, GNorm, Gpu, Matmul, Strides};
|
||||
use crate::memattn::{GRID, rope_tables};
|
||||
use crate::weights::Weights;
|
||||
use crate::{Error, Result};
|
||||
|
||||
const DIM: usize = 256;
|
||||
const LAYERS: usize = 4;
|
||||
|
||||
struct GAttn {
|
||||
q: GLinear,
|
||||
k: GLinear,
|
||||
v: GLinear,
|
||||
o: GLinear,
|
||||
}
|
||||
|
||||
struct GLayer {
|
||||
self_attn: GAttn,
|
||||
cross: GAttn,
|
||||
ln1: GNorm,
|
||||
ln2: GNorm,
|
||||
ln3: GNorm,
|
||||
lin1: GLinear,
|
||||
lin2: GLinear,
|
||||
}
|
||||
|
||||
pub struct GpuMemoryAttention {
|
||||
gpu: Arc<Gpu>,
|
||||
layers: Vec<GLayer>,
|
||||
final_ln: GNorm,
|
||||
cos: Buf,
|
||||
sin: Buf,
|
||||
}
|
||||
|
||||
fn attn(g: &Gpu, w: &Weights, prefix: &str, kv: usize) -> Result<GAttn> {
|
||||
Ok(GAttn {
|
||||
q: GLinear::load(g, w, &format!("{prefix}.q_proj"), DIM, DIM)?,
|
||||
k: GLinear::load(g, w, &format!("{prefix}.k_proj"), kv, DIM)?,
|
||||
v: GLinear::load(g, w, &format!("{prefix}.v_proj"), kv, DIM)?,
|
||||
o: GLinear::load(g, w, &format!("{prefix}.o_proj"), DIM, DIM)?,
|
||||
})
|
||||
}
|
||||
|
||||
impl GpuMemoryAttention {
|
||||
pub fn load(gpu: Arc<Gpu>, w: &Weights) -> Result<GpuMemoryAttention> {
|
||||
let layers = (0..LAYERS)
|
||||
.map(|i| {
|
||||
let p = format!("memory_attention.layers.{i}");
|
||||
Ok(GLayer {
|
||||
self_attn: attn(&gpu, w, &format!("{p}.self_attn"), DIM)?,
|
||||
cross: attn(&gpu, w, &format!("{p}.cross_attn_image"), 64)?,
|
||||
ln1: GNorm::load(&gpu, w, &format!("{p}.layer_norm1"), DIM)?,
|
||||
ln2: GNorm::load(&gpu, w, &format!("{p}.layer_norm2"), DIM)?,
|
||||
ln3: GNorm::load(&gpu, w, &format!("{p}.layer_norm3"), DIM)?,
|
||||
lin1: GLinear::load(&gpu, w, &format!("{p}.linear1"), DIM, 2048)?,
|
||||
lin2: GLinear::load(&gpu, w, &format!("{p}.linear2"), 2048, DIM)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let (cos, sin) = rope_tables(10000.0)?;
|
||||
let flat = |t: Tensor| -> Result<Vec<f32>> { Ok(t.flatten_all()?.to_vec1::<f32>()?) };
|
||||
Ok(GpuMemoryAttention {
|
||||
final_ln: GNorm::load(&gpu, w, "memory_attention.layer_norm", DIM)?,
|
||||
cos: gpu.upload(&flat(cos)?),
|
||||
sin: gpu.upload(&flat(sin)?),
|
||||
layers,
|
||||
gpu,
|
||||
})
|
||||
}
|
||||
|
||||
fn norm(&self, x: &Buf, rows: usize, n: &GNorm) -> Result<Buf> {
|
||||
let out = self.gpu.zeros(rows * DIM)?;
|
||||
self.gpu.layernorm(x, &n.w, &n.b, &out, rows, DIM, 1e-5, false)?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn rope(&self, x: &Buf, rows: usize, rotated: usize) -> Result<Buf> {
|
||||
let out = self.gpu.zeros(rows * DIM)?;
|
||||
self.gpu.rope(x, &self.cos, &self.sin, &out, rows, DIM, GRID * GRID, rotated)?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// `softmax(q . k^T / 16) . v` for `[nq, 256]` queries and `[nk, 256]` keys and values, through `scores`
|
||||
/// (at least `nq * nk` elements); returns `[nq, 256]`.
|
||||
fn attention(&self, q: &Buf, k: &Buf, v: &Buf, nq: usize, nk: usize, scores: &Buf) -> Result<Buf> {
|
||||
let g = &self.gpu;
|
||||
g.matmul(&Matmul {
|
||||
m: nq,
|
||||
n: nk,
|
||||
k: DIM,
|
||||
batch: 1,
|
||||
a: (q, Strides::row_major(DIM)),
|
||||
b: (k, Strides { offset: 0, batch: 0, row: 1, col: DIM }),
|
||||
bias: None,
|
||||
c: (scores, Strides::row_major(nk)),
|
||||
residual: None,
|
||||
alpha: (DIM as f32).powf(-0.5),
|
||||
act: Act::None,
|
||||
})?;
|
||||
g.softmax(scores, nq, nk)?;
|
||||
let out = g.zeros(nq * DIM)?;
|
||||
g.matmul(&Matmul {
|
||||
m: nq,
|
||||
n: DIM,
|
||||
k: nk,
|
||||
batch: 1,
|
||||
a: (scores, Strides::row_major(nk)),
|
||||
b: (v, Strides::row_major(DIM)),
|
||||
bias: None,
|
||||
c: (&out, Strides::row_major(DIM)),
|
||||
residual: None,
|
||||
alpha: 1.0,
|
||||
act: Act::None,
|
||||
})?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// As [`crate::memattn::MemoryAttention::forward`]: `current` and `current_pos` `[4096, 256]`, `memory` and
|
||||
/// `memory_pos` `[M, 64]`, the last `num_pointers` memory rows being object pointers (not rotated).
|
||||
pub fn forward(&self, current: &Tensor, current_pos: &Tensor, memory: &Tensor, memory_pos: &Tensor, num_pointers: usize) -> Result<Tensor> {
|
||||
let g = &self.gpu;
|
||||
let _turn = g.exclusive();
|
||||
let n = GRID * GRID;
|
||||
if current.dims() != [n, DIM] || current_pos.dims() != [n, DIM] {
|
||||
return Err(Error::Input(format!("current features must be [{n}, {DIM}], got {:?}", current.dims())));
|
||||
}
|
||||
let (m, md) = memory.dims2()?;
|
||||
if md != 64 || memory_pos.dims() != [m, 64] {
|
||||
return Err(Error::Input(format!("memory must be [M, 64] with a matching position encoding, got {:?}", memory.dims())));
|
||||
}
|
||||
let rotated = m.checked_sub(num_pointers).ok_or_else(|| Error::Input("more object pointers than memory tokens".into()))?;
|
||||
if rotated % n != 0 {
|
||||
return Err(Error::Input(format!("memory length {rotated} is not a multiple of {n}")));
|
||||
}
|
||||
let flat = |t: &Tensor| -> Result<Vec<f32>> { Ok(t.contiguous()?.flatten_all()?.to_vec1::<f32>()?) };
|
||||
let mut x = g.upload(&flat(&(current + (current_pos * 0.1)?)?)?);
|
||||
let mem = g.upload(&flat(memory)?);
|
||||
let keys_in = g.upload(&flat(&(memory + memory_pos)?)?);
|
||||
let scores = g.zeros(n * n.max(m))?;
|
||||
for l in &self.layers {
|
||||
// self-attention over the frame
|
||||
let q_in = self.norm(&x, n, &l.ln1)?;
|
||||
let (q, k, v) = (g.zeros(n * DIM)?, g.zeros(n * DIM)?, g.zeros(n * DIM)?);
|
||||
g.linear(&q_in, n, &l.self_attn.q, &q, Act::None, None)?;
|
||||
g.linear(&q_in, n, &l.self_attn.k, &k, Act::None, None)?;
|
||||
g.linear(&q_in, n, &l.self_attn.v, &v, Act::None, None)?;
|
||||
let (q, k) = (self.rope(&q, n, n)?, self.rope(&k, n, n)?);
|
||||
let a = self.attention(&q, &k, &v, n, n, &scores)?;
|
||||
let x1 = g.zeros(n * DIM)?;
|
||||
g.linear(&a, n, &l.self_attn.o, &x1, Act::None, Some(&x))?;
|
||||
// cross-attention to the memory bank
|
||||
let q_in = self.norm(&x1, n, &l.ln2)?;
|
||||
let q = g.zeros(n * DIM)?;
|
||||
g.linear(&q_in, n, &l.cross.q, &q, Act::None, None)?;
|
||||
let q = self.rope(&q, n, n)?;
|
||||
let (k, v) = (g.zeros(m * DIM)?, g.zeros(m * DIM)?);
|
||||
g.linear(&keys_in, m, &l.cross.k, &k, Act::None, None)?;
|
||||
g.linear(&mem, m, &l.cross.v, &v, Act::None, None)?;
|
||||
let k = self.rope(&k, m, rotated)?;
|
||||
let a = self.attention(&q, &k, &v, n, m, &scores)?;
|
||||
let x2 = g.zeros(n * DIM)?;
|
||||
g.linear(&a, n, &l.cross.o, &x2, Act::None, Some(&x1))?;
|
||||
// feed-forward
|
||||
let f_in = self.norm(&x2, n, &l.ln3)?;
|
||||
let hidden = g.zeros(n * 2048)?;
|
||||
g.linear(&f_in, n, &l.lin1, &hidden, Act::Relu, None)?;
|
||||
let x3 = g.zeros(n * DIM)?;
|
||||
g.linear(&hidden, n, &l.lin2, &x3, Act::None, Some(&x2))?;
|
||||
x = x3;
|
||||
}
|
||||
let out = self.norm(&x, n, &self.final_ln)?;
|
||||
g.flush()?;
|
||||
Ok(Tensor::from_vec(g.download(&out)?, (n, DIM), &Device::Cpu)?)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
// Original work for EffectCraft; the maths follows the Hugging Face `transformers` SAM 2 video
|
||||
// implementation (Apache-2.0, see NOTICE).
|
||||
//! The memory encoder on the GPU: the mask is down-sampled by four strided convolutions, fused with the
|
||||
//! frame's features by two ConvNeXt-style blocks, and projected to the 64 channels of a memory.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use candle_core::{Device, Tensor};
|
||||
|
||||
use super::{Act, Buf, GLinear, GNorm, Gpu};
|
||||
use crate::weights::Weights;
|
||||
use crate::{Error, Result};
|
||||
|
||||
struct GConv {
|
||||
w: Buf,
|
||||
b: Buf,
|
||||
ci: usize,
|
||||
co: usize,
|
||||
k: usize,
|
||||
}
|
||||
|
||||
struct GFuser {
|
||||
dw: GConv,
|
||||
norm: GNorm,
|
||||
pw1: GLinear,
|
||||
/// The layer scale is folded into this layer's weights.
|
||||
pw2: GLinear,
|
||||
}
|
||||
|
||||
pub struct GpuMemoryEncoder {
|
||||
gpu: Arc<Gpu>,
|
||||
down: Vec<(GConv, GNorm)>,
|
||||
final_conv: GLinear,
|
||||
feature_projection: GLinear,
|
||||
fuser: Vec<GFuser>,
|
||||
projection: GLinear,
|
||||
}
|
||||
|
||||
fn conv(g: &Gpu, w: &Weights, name: &str, ci: usize, co: usize, k: usize, groups: usize) -> Result<GConv> {
|
||||
let flat = |t: Tensor| -> Result<Vec<f32>> { Ok(t.flatten_all()?.to_vec1::<f32>()?) };
|
||||
Ok(GConv {
|
||||
w: g.upload(&flat(w.get(&format!("{name}.weight"), &[co, ci / groups, k, k])?)?),
|
||||
b: g.upload(&flat(w.get(&format!("{name}.bias"), &[co])?)?),
|
||||
ci,
|
||||
co,
|
||||
k,
|
||||
})
|
||||
}
|
||||
|
||||
impl GpuMemoryEncoder {
|
||||
pub fn load(gpu: Arc<Gpu>, w: &Weights) -> Result<GpuMemoryEncoder> {
|
||||
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(&gpu, w, &format!("{p}.conv"), ch, out, 3, 1)?, GNorm::load(&gpu, w, &format!("{p}.layer_norm"), out)?));
|
||||
ch = out;
|
||||
}
|
||||
let fuser = (0..2)
|
||||
.map(|i| {
|
||||
let p = format!("memory_encoder.memory_fuser.layers.{i}");
|
||||
let scale = w.get(&format!("{p}.scale"), &[256])?.flatten_all()?.to_vec1::<f32>()?;
|
||||
Ok(GFuser {
|
||||
dw: conv(&gpu, w, &format!("{p}.depthwise_conv"), 256, 256, 7, 256)?,
|
||||
norm: GNorm::load(&gpu, w, &format!("{p}.layer_norm"), 256)?,
|
||||
pw1: GLinear::load(&gpu, w, &format!("{p}.pointwise_conv1"), 256, 1024)?,
|
||||
pw2: GLinear::load_scaled(&gpu, w, &format!("{p}.pointwise_conv2"), 1024, 256, &scale)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
Ok(GpuMemoryEncoder {
|
||||
final_conv: GLinear::load_conv1x1(&gpu, w, "memory_encoder.mask_downsampler.final_conv", 256, 256)?,
|
||||
feature_projection: GLinear::load_conv1x1(&gpu, w, "memory_encoder.feature_projection", 256, 256)?,
|
||||
projection: GLinear::load_conv1x1(&gpu, w, "memory_encoder.projection", 256, 64)?,
|
||||
down,
|
||||
fuser,
|
||||
gpu,
|
||||
})
|
||||
}
|
||||
|
||||
/// `features`: `[1, 256, 64, 64]`; `mask`: `[1, 1, 1024, 1024]` already scaled (`20 m - 10`). Returns the raw
|
||||
/// memory tokens `[4096, 64]` (the occlusion embedding and bfloat16 rounding are the caller's).
|
||||
pub fn encode(&self, features: &Tensor, mask: &Tensor) -> Result<Tensor> {
|
||||
let g = &self.gpu;
|
||||
let _turn = g.exclusive();
|
||||
let flat = |t: &Tensor| -> Result<Vec<f32>> { Ok(t.contiguous()?.flatten_all()?.to_vec1::<f32>()?) };
|
||||
if features.dims() != [1, 256, 64, 64] || mask.dims() != [1, 1, 1024, 1024] {
|
||||
return Err(Error::Input(format!("memory encoder input shapes {:?} / {:?}", features.dims(), mask.dims())));
|
||||
}
|
||||
let tokens = features.permute((0, 2, 3, 1))?.contiguous()?;
|
||||
let ftok = g.upload(&flat(&tokens)?);
|
||||
let mut cur = g.upload(&flat(mask)?);
|
||||
let (mut h, mut w) = (1024usize, 1024usize);
|
||||
for (c, ln) in &self.down {
|
||||
let (ho, wo) = (h / 2, w / 2);
|
||||
let raw = g.zeros(ho * wo * c.co)?;
|
||||
g.conv2d(&cur, &c.w, &c.b, &raw, (h, w, c.ci, c.co), c.k, 2, 1, 1)?;
|
||||
let normed = g.zeros(ho * wo * c.co)?;
|
||||
g.layernorm(&raw, &ln.w, &ln.b, &normed, ho * wo, c.co, 1e-6, true)?;
|
||||
(cur, h, w) = (normed, ho, wo);
|
||||
}
|
||||
let n = h * w;
|
||||
if n != 4096 {
|
||||
return Err(Error::Compute(format!("the mask down-sampler ended at {w}x{h}, not 64x64")));
|
||||
}
|
||||
let m = g.zeros(n * 256)?;
|
||||
g.linear(&cur, n, &self.final_conv, &m, Act::None, None)?;
|
||||
let mut x = g.zeros(n * 256)?;
|
||||
g.linear(&ftok, n, &self.feature_projection, &x, Act::None, Some(&m))?;
|
||||
for b in &self.fuser {
|
||||
let d = g.zeros(n * 256)?;
|
||||
g.conv2d(&x, &b.dw.w, &b.dw.b, &d, (h, w, 256, 256), 7, 1, 3, 256)?;
|
||||
let normed = g.zeros(n * 256)?;
|
||||
g.layernorm(&d, &b.norm.w, &b.norm.b, &normed, n, 256, 1e-6, false)?;
|
||||
let hidden = g.zeros(n * 1024)?;
|
||||
g.linear(&normed, n, &b.pw1, &hidden, Act::Gelu, None)?;
|
||||
let next = g.zeros(n * 256)?;
|
||||
g.linear(&hidden, n, &b.pw2, &next, Act::None, Some(&x))?;
|
||||
x = next;
|
||||
}
|
||||
let out = g.zeros(n * 64)?;
|
||||
g.linear(&x, n, &self.projection, &out, Act::None, None)?;
|
||||
g.flush()?;
|
||||
Ok(Tensor::from_vec(g.download(&out)?, (n, 64), &Device::Cpu)?)
|
||||
}
|
||||
}
|
||||
@@ -72,6 +72,7 @@ pub struct Gpu {
|
||||
module: wgpu::ShaderModule,
|
||||
pipelines: Mutex<HashMap<&'static str, Arc<wgpu::ComputePipeline>>>,
|
||||
encoder: Mutex<Option<wgpu::CommandEncoder>>,
|
||||
turn: Mutex<()>,
|
||||
error: Arc<Mutex<Option<String>>>,
|
||||
adapter: String,
|
||||
}
|
||||
@@ -89,8 +90,16 @@ fn u32_of(v: usize, what: &str) -> Result<u32> {
|
||||
}
|
||||
|
||||
impl Gpu {
|
||||
/// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders
|
||||
/// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go,
|
||||
/// the device stays. An error (no GPU) is remembered too.
|
||||
pub fn shared() -> Result<Arc<Gpu>> {
|
||||
static DEVICE: std::sync::OnceLock<std::result::Result<Arc<Gpu>, String>> = std::sync::OnceLock::new();
|
||||
DEVICE.get_or_init(|| Gpu::new().map(Arc::new).map_err(|e| e.to_string())).clone().map_err(Error::Compute)
|
||||
}
|
||||
|
||||
/// A device on the best real GPU (a software adapter is refused: the CPU path is faster).
|
||||
pub fn new() -> Result<Gpu> {
|
||||
fn new() -> Result<Gpu> {
|
||||
let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
|
||||
let adapter = pollster::block_on(
|
||||
instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }),
|
||||
@@ -117,7 +126,13 @@ impl Gpu {
|
||||
sink.lock().unwrap_or_else(PoisonError::into_inner).get_or_insert_with(|| format!("GPU device lost ({reason:?}): {message}"));
|
||||
});
|
||||
let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("sam2 kernels"), source: wgpu::ShaderSource::Wgsl(SHADER.into()) });
|
||||
Ok(Gpu { device, queue, module, pipelines: Mutex::default(), encoder: Mutex::new(None), error, adapter: info.name })
|
||||
Ok(Gpu { device, queue, module, pipelines: Mutex::default(), encoder: Mutex::new(None), turn: Mutex::new(()), error, adapter: info.name })
|
||||
}
|
||||
|
||||
/// One GPU pass at a time: a caller holds this while it queues, flushes and reads back a whole job, so
|
||||
/// two threads never interleave their dispatches in one command stream.
|
||||
pub fn exclusive(&self) -> std::sync::MutexGuard<'_, ()> {
|
||||
self.turn.lock().unwrap_or_else(PoisonError::into_inner)
|
||||
}
|
||||
|
||||
pub fn adapter_name(&self) -> &str {
|
||||
@@ -304,6 +319,8 @@ impl Gpu {
|
||||
}
|
||||
|
||||
pub mod encoder;
|
||||
pub mod memattn;
|
||||
pub mod memenc;
|
||||
mod ops;
|
||||
|
||||
/// A linear layer's weight `[out, in]` and bias `[out]` on the GPU.
|
||||
@@ -335,6 +352,24 @@ impl GLinear {
|
||||
})
|
||||
}
|
||||
|
||||
/// As [`GLinear::load`], with output channel `j` scaled by `scale[j]` (the memory fuser's layer scale).
|
||||
pub fn load_scaled(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize, scale: &[f32]) -> Result<GLinear> {
|
||||
if scale.len() != out {
|
||||
return Err(Error::Weights(format!("`{name}`: {} scales for {out} outputs", scale.len())));
|
||||
}
|
||||
let mut wt = flat(w.get(&format!("{name}.weight"), &[out, inp])?)?;
|
||||
let mut bias = flat(w.get(&format!("{name}.bias"), &[out])?)?;
|
||||
for (j, s) in scale.iter().enumerate() {
|
||||
if let Some(row) = wt.get_mut(j * inp..(j + 1) * inp) {
|
||||
row.iter_mut().for_each(|v| *v *= s);
|
||||
}
|
||||
if let Some(b) = bias.get_mut(j) {
|
||||
*b *= s;
|
||||
}
|
||||
}
|
||||
Ok(GLinear { w: g.upload(&wt), b: g.upload(&bias), inp, out })
|
||||
}
|
||||
|
||||
/// A 1 x 1 convolution (`[out, in, 1, 1]`) as a linear layer over channels-last tokens.
|
||||
pub fn load_conv1x1(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result<GLinear> {
|
||||
Ok(GLinear {
|
||||
|
||||
@@ -27,12 +27,18 @@ fn need(buf: &Buf, len: usize, what: &str) -> Result<()> {
|
||||
|
||||
impl Gpu {
|
||||
/// `y = layer_norm(x) * weight + bias` over the last dimension of a `[rows, cols]` buffer.
|
||||
pub fn layernorm(&self, x: &Buf, weight: &Buf, bias: &Buf, y: &Buf, rows: usize, cols: usize, eps: f32) -> Result<()> {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn layernorm(&self, x: &Buf, weight: &Buf, bias: &Buf, y: &Buf, rows: usize, cols: usize, eps: f32, gelu: bool) -> Result<()> {
|
||||
need(x, rows * cols, "layernorm x")?;
|
||||
need(y, rows * cols, "layernorm y")?;
|
||||
need(weight, cols, "layernorm weight")?;
|
||||
need(bias, cols, "layernorm bias")?;
|
||||
self.run("layernorm", &[(0, x), (1, weight), (2, bias), (3, y)], &[n32(rows, "rows")?, n32(cols, "cols")?, eps.to_bits()], grid(rows)?);
|
||||
self.run(
|
||||
"layernorm",
|
||||
&[(0, x), (1, weight), (2, bias), (3, y)],
|
||||
&[n32(rows, "rows")?, n32(cols, "cols")?, eps.to_bits(), u32::from(gelu)],
|
||||
grid(rows)?,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -43,6 +49,52 @@ impl Gpu {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A convolution over channels-last `[h, w, ci]` with weight `[co, ci / groups, k, k]` and bias `[co]`,
|
||||
/// output `[ho, wo, co]` (`groups == ci == co` is depthwise).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn conv2d(
|
||||
&self,
|
||||
x: &Buf,
|
||||
weight: &Buf,
|
||||
bias: &Buf,
|
||||
out: &Buf,
|
||||
dims: (usize, usize, usize, usize),
|
||||
k: usize,
|
||||
stride: usize,
|
||||
pad: usize,
|
||||
groups: usize,
|
||||
) -> Result<()> {
|
||||
let (h, w, ci, co) = dims;
|
||||
if k == 0 || stride == 0 || groups == 0 || ci % groups != 0 || co % groups != 0 || h + 2 * pad < k || w + 2 * pad < k {
|
||||
return Err(Error::Input(format!("a {k}x{k} convolution (stride {stride}, groups {groups}) does not fit {w}x{h}x{ci} -> {co}")));
|
||||
}
|
||||
let (ho, wo) = ((h + 2 * pad - k) / stride + 1, (w + 2 * pad - k) / stride + 1);
|
||||
need(x, h * w * ci, "conv2d input")?;
|
||||
need(weight, co * (ci / groups) * k * k, "conv2d weight")?;
|
||||
need(bias, co, "conv2d bias")?;
|
||||
need(out, ho * wo * co, "conv2d output")?;
|
||||
let n = ho * wo * co;
|
||||
self.run(
|
||||
"conv2d",
|
||||
&[(0, x), (1, weight), (2, bias), (3, out)],
|
||||
&[
|
||||
n32(n, "n")?,
|
||||
n32(h, "h")?,
|
||||
n32(w, "w")?,
|
||||
n32(ci, "ci")?,
|
||||
n32(co, "co")?,
|
||||
n32(ho, "ho")?,
|
||||
n32(wo, "wo")?,
|
||||
n32(k, "k")?,
|
||||
n32(stride, "stride")?,
|
||||
n32(pad, "pad")?,
|
||||
n32(groups, "groups")?,
|
||||
],
|
||||
items(n)?,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `out = a + b` over `n` elements.
|
||||
pub fn add(&self, a: &Buf, b: &Buf, out: &Buf, n: usize) -> Result<()> {
|
||||
need(a, n, "add a")?;
|
||||
|
||||
@@ -6,27 +6,25 @@ use super::*;
|
||||
|
||||
/// The one device the tests share (several devices at once upset the Vulkan loader), held for the whole
|
||||
/// test so the tests take turns.
|
||||
pub(crate) struct Shared(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>, &'static Gpu);
|
||||
pub(crate) struct Shared(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>, Arc<Gpu>);
|
||||
|
||||
impl std::ops::Deref for Shared {
|
||||
type Target = Gpu;
|
||||
fn deref(&self) -> &Gpu {
|
||||
self.1
|
||||
&self.1
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn gpu() -> Option<Shared> {
|
||||
static TURN: Mutex<()> = Mutex::new(());
|
||||
static DEVICE: std::sync::OnceLock<Option<Gpu>> = std::sync::OnceLock::new();
|
||||
let guard = TURN.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
let device = DEVICE.get_or_init(|| match Gpu::new() {
|
||||
Ok(g) => Some(g),
|
||||
match Gpu::shared() {
|
||||
Ok(g) => Some(Shared(guard, g)),
|
||||
Err(e) => {
|
||||
eprintln!("skipped: {e}");
|
||||
None
|
||||
}
|
||||
});
|
||||
device.as_ref().map(|g| Shared(guard, g))
|
||||
}
|
||||
}
|
||||
|
||||
/// Deterministic pseudo-random values in -1..1.
|
||||
@@ -176,7 +174,7 @@ fn layernorm_and_softmax_match_the_definition() {
|
||||
let x = noise(rows * cols, 7);
|
||||
let (w, b) = (noise(cols, 8), noise(cols, 9));
|
||||
let (gx, gw, gb, gy) = (g.upload(&x), g.upload(&w), g.upload(&b), g.zeros(rows * cols).unwrap());
|
||||
g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6).unwrap();
|
||||
g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6, false).unwrap();
|
||||
let want: Vec<f32> = x
|
||||
.chunks(cols)
|
||||
.flat_map(|r| {
|
||||
@@ -299,3 +297,50 @@ fn patch_embed_rope_and_upsample_match_the_cpu_definitions() {
|
||||
let want: Vec<f32> = (0..h * w * c).map(|i| lat[i] + prev[((i / c / w / 2) * (w / 2) + (i / c % w) / 2) * c + i % c]).collect();
|
||||
assert_close(&run_and_read(&g, &go), &want, 0.0, "upsample_add");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convolutions_and_the_layernorm_gelu_epilogue_match_candle() {
|
||||
use candle_core::{Device, Tensor};
|
||||
let Some(g) = gpu() else { return };
|
||||
// (h, w, ci, co, k, stride, pad, groups): a strided 3 x 3, a depthwise 7 x 7, and a 1-channel input
|
||||
for (h, w, ci, co, k, stride, pad, groups) in
|
||||
[(9usize, 13usize, 3usize, 5usize, 3usize, 2usize, 1usize, 1usize), (10, 10, 4, 4, 7, 1, 3, 4), (16, 16, 1, 4, 3, 2, 1, 1)]
|
||||
{
|
||||
let (x, wt, b) = (noise(h * w * ci, 30), noise(co * (ci / groups) * k * k, 31), noise(co, 32));
|
||||
let (gx, gw, gb) = (g.upload(&x), g.upload(&wt), g.upload(&b));
|
||||
let (ho, wo) = ((h + 2 * pad - k) / stride + 1, (w + 2 * pad - k) / stride + 1);
|
||||
let out = g.zeros(ho * wo * co).unwrap();
|
||||
g.conv2d(&gx, &gw, &gb, &out, (h, w, ci, co), k, stride, pad, groups).unwrap();
|
||||
// candle works on NCHW: bring the channels-last input over, convolve, bring the result back
|
||||
let t = Tensor::from_vec(x, (1, h, w, ci), &Device::Cpu).unwrap().permute((0, 3, 1, 2)).unwrap().contiguous().unwrap();
|
||||
let kern = Tensor::from_vec(wt, (co, ci / groups, k, k), &Device::Cpu).unwrap();
|
||||
let want = t
|
||||
.conv2d(&kern, pad, stride, 1, groups)
|
||||
.unwrap()
|
||||
.broadcast_add(&Tensor::from_vec(b, (1, co, 1, 1), &Device::Cpu).unwrap())
|
||||
.unwrap()
|
||||
.permute((0, 2, 3, 1))
|
||||
.unwrap();
|
||||
assert_close(
|
||||
&run_and_read(&g, &out),
|
||||
&want.flatten_all().unwrap().to_vec1::<f32>().unwrap(),
|
||||
2e-4,
|
||||
&format!("conv {k}x{k} stride {stride} groups {groups}"),
|
||||
);
|
||||
}
|
||||
assert!(g.conv2d(&g.upload(&[0.0; 4]), &g.upload(&[0.0; 4]), &g.upload(&[0.0; 1]), &g.zeros(4).unwrap(), (2, 2, 1, 1), 5, 1, 0, 1).is_err());
|
||||
// layer norm followed by GELU in the same pass
|
||||
let (rows, cols) = (50usize, 96usize);
|
||||
let (x, wt, b) = (noise(rows * cols, 33), noise(cols, 34), noise(cols, 35));
|
||||
let (gx, gw, gb, gy) = (g.upload(&x), g.upload(&wt), g.upload(&b), g.zeros(rows * cols).unwrap());
|
||||
g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6, true).unwrap();
|
||||
let want: Vec<f32> = x
|
||||
.chunks(cols)
|
||||
.flat_map(|r| {
|
||||
let mean = r.iter().sum::<f32>() / cols as f32;
|
||||
let var = r.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / cols as f32;
|
||||
r.iter().enumerate().map(|(i, v)| gelu((v - mean) / (var + 1e-6).sqrt() * wt[i] + b[i])).collect::<Vec<_>>()
|
||||
})
|
||||
.collect();
|
||||
assert_close(&run_and_read(&g, &gy), &want, 1e-3, "layernorm + gelu");
|
||||
}
|
||||
|
||||
@@ -79,7 +79,7 @@ impl Sam2Mask {
|
||||
{
|
||||
return Ok(e.clone());
|
||||
}
|
||||
let enc = Arc::new(self.model.image.encode(&preprocess(rgb, w, h)?)?);
|
||||
let enc = Arc::new(self.model.encode(&preprocess(rgb, w, h)?)?);
|
||||
*self.last.lock().unwrap_or_else(PoisonError::into_inner) = Some((key, enc.clone()));
|
||||
Ok(enc)
|
||||
}
|
||||
@@ -117,6 +117,10 @@ impl MaskModel for Sam2Mask {
|
||||
Some(self)
|
||||
}
|
||||
|
||||
fn device_name(&self) -> Option<String> {
|
||||
Some(self.model.device_name())
|
||||
}
|
||||
|
||||
fn idle_secs(&self) -> Option<f64> {
|
||||
Some(self.used.lock().unwrap_or_else(PoisonError::into_inner).elapsed().as_secs_f64())
|
||||
}
|
||||
@@ -129,7 +133,7 @@ pub fn load(bytes: &[u8]) -> std::result::Result<Arc<Sam2Mask>, String> {
|
||||
return Err(format!("{}: the file does not match (SHA-256 {got}, expected {})", SAM2_TINY.name, SAM2_TINY.sha256));
|
||||
}
|
||||
let w = crate::weights::Weights::from_bytes(bytes).map_err(|e| e.to_string())?;
|
||||
let model = VideoModel::load(&w).map_err(|e| e.to_string())?;
|
||||
let model = VideoModel::load_auto(&w).map_err(|e| e.to_string())?;
|
||||
Ok(Arc::new(Sam2Mask::new(Arc::new(model))))
|
||||
}
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ const LAYERS: usize = 4;
|
||||
pub const GRID: usize = 64;
|
||||
|
||||
/// `cos` and `sin` of the axial rotary encoding for the `GRID x GRID` grid, each `[GRID², 256]`.
|
||||
fn rope_tables(theta: f32) -> Result<(Tensor, Tensor)> {
|
||||
pub(crate) fn rope_tables(theta: f32) -> Result<(Tensor, Tensor)> {
|
||||
let freqs = DIM / 2 / 2; // per axis: head_dim / 4
|
||||
let inv: Vec<f32> = (0..freqs).map(|i| 1.0 / theta.powf(i as f32 / freqs as f32)).collect();
|
||||
let n = GRID * GRID;
|
||||
|
||||
@@ -123,10 +123,16 @@ impl MemoryEncoder {
|
||||
/// 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)? })
|
||||
let (x, _) = self.encode(features, mask)?;
|
||||
self.finish(&x.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?, object_present)
|
||||
}
|
||||
|
||||
/// A memory from the encoder's raw tokens `[HW, 64]`: the occlusion embedding when the object is absent,
|
||||
/// rounded through bfloat16, with the position encoding (shared by the CPU and GPU encoders).
|
||||
pub fn finish(&self, tokens: &Tensor, object_present: bool) -> Result<Memory> {
|
||||
let tokens = if object_present { tokens.clone() } else { tokens.broadcast_add(&self.occlusion)? };
|
||||
let features = tokens.to_dtype(candle_core::DType::BF16)?.to_dtype(candle_core::DType::F32)?;
|
||||
let pos = sine_position(32, 64, 64)?;
|
||||
Ok(Memory { features, pos: pos.flatten_from(2)?.squeeze(0)?.t()?.contiguous()? })
|
||||
}
|
||||
}
|
||||
|
||||
+149
-6
@@ -14,9 +14,14 @@ use candle_core::{DType, Tensor};
|
||||
use candle_nn::{Linear, ops::sigmoid};
|
||||
|
||||
use crate::decoder::FeedForward;
|
||||
use crate::gpu::Gpu;
|
||||
use crate::gpu::encoder::GpuEncoder;
|
||||
use crate::gpu::memattn::GpuMemoryAttention;
|
||||
use crate::gpu::memenc::GpuMemoryEncoder;
|
||||
use crate::memattn::MemoryAttention;
|
||||
use crate::memory::{Memory, MemoryEncoder};
|
||||
use crate::model::ImageModel;
|
||||
use crate::model::{Encoded, ImageModel};
|
||||
use crate::neck::ImageFeatures;
|
||||
use crate::nn::{apply, linear};
|
||||
use crate::prompt::{INPUT_SIZE, Prompt};
|
||||
use crate::resize::antialias_bilinear;
|
||||
@@ -79,8 +84,31 @@ pub struct FrameMask {
|
||||
pub occluded: bool,
|
||||
}
|
||||
|
||||
/// Which device runs the heavy parts (the backbone, the neck and the memory attention).
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Backend {
|
||||
/// Candle on the CPU: the reference path.
|
||||
Cpu,
|
||||
/// wgpu compute on the GPU (any vendor); an error if there is none.
|
||||
Gpu,
|
||||
/// The GPU when there is a real one and it loads, else the CPU. The environment variable
|
||||
/// `EFFECTCRAFT_SAM2_DEVICE` (`cpu`, `gpu` or `auto`) overrides the choice.
|
||||
Auto,
|
||||
}
|
||||
|
||||
/// The GPU path: kernels for the backbone and the memory attention. Any failure switches it off for good
|
||||
/// and the CPU path takes over.
|
||||
struct GpuParts {
|
||||
encoder: GpuEncoder,
|
||||
attention: GpuMemoryAttention,
|
||||
memory_encoder: GpuMemoryEncoder,
|
||||
ok: AtomicBool,
|
||||
adapter: String,
|
||||
}
|
||||
|
||||
/// The weights of everything the video tracker needs.
|
||||
pub struct VideoModel {
|
||||
gpu: Option<GpuParts>,
|
||||
pub image: ImageModel,
|
||||
memory_encoder: MemoryEncoder,
|
||||
memory_attention: MemoryAttention,
|
||||
@@ -93,9 +121,37 @@ pub struct VideoModel {
|
||||
}
|
||||
|
||||
impl VideoModel {
|
||||
/// On the CPU: the reference path every test compares against.
|
||||
pub fn load(w: &Weights) -> Result<VideoModel> {
|
||||
VideoModel::load_with(w, Backend::Cpu)
|
||||
}
|
||||
|
||||
/// On the GPU when there is one (see [`Backend::Auto`]).
|
||||
pub fn load_auto(w: &Weights) -> Result<VideoModel> {
|
||||
VideoModel::load_with(w, Backend::Auto)
|
||||
}
|
||||
|
||||
pub fn load_with(w: &Weights, backend: Backend) -> Result<VideoModel> {
|
||||
let backend = match (backend, std::env::var("EFFECTCRAFT_SAM2_DEVICE").ok().as_deref()) {
|
||||
(Backend::Auto, Some("cpu")) => Backend::Cpu,
|
||||
(Backend::Auto, Some("gpu")) => Backend::Gpu,
|
||||
(b, _) => b,
|
||||
};
|
||||
let image = ImageModel::load(w)?;
|
||||
let gpu = match backend {
|
||||
Backend::Cpu => None,
|
||||
Backend::Gpu => Some(Self::load_gpu(w, &image)?),
|
||||
Backend::Auto => match Self::load_gpu(w, &image) {
|
||||
Ok(g) => Some(g),
|
||||
Err(e) => {
|
||||
log::info!("SAM 2 runs on the CPU: {e}");
|
||||
None
|
||||
}
|
||||
},
|
||||
};
|
||||
Ok(VideoModel {
|
||||
image: ImageModel::load(w)?,
|
||||
gpu,
|
||||
image,
|
||||
memory_encoder: MemoryEncoder::load(w)?,
|
||||
memory_attention: MemoryAttention::load(w)?,
|
||||
temporal_pos: w.get("memory_temporal_positional_encoding", &[MEMORY_FRAMES, 1, 1, 64])?.flatten_all()?.to_vec1::<f32>()?,
|
||||
@@ -106,6 +162,79 @@ impl VideoModel {
|
||||
})
|
||||
}
|
||||
|
||||
fn load_gpu(w: &Weights, image: &ImageModel) -> Result<GpuParts> {
|
||||
let gpu = Gpu::shared()?;
|
||||
Ok(GpuParts {
|
||||
encoder: GpuEncoder::load(gpu.clone(), w, image.vision.hiera.clone())?,
|
||||
attention: GpuMemoryAttention::load(gpu.clone(), w)?,
|
||||
memory_encoder: GpuMemoryEncoder::load(gpu.clone(), w)?,
|
||||
ok: AtomicBool::new(true),
|
||||
adapter: gpu.adapter_name().to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// The device the heavy parts run on now: the GPU's name, or "CPU".
|
||||
pub fn device_name(&self) -> String {
|
||||
match &self.gpu {
|
||||
Some(g) if g.ok.load(Ordering::Relaxed) => g.adapter.clone(),
|
||||
_ => "CPU".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Run on the CPU from now on (a GPU that misbehaves, or the user's choice). Weights and results are
|
||||
/// unchanged; only the device differs.
|
||||
pub fn disable_gpu(&self) {
|
||||
if let Some(g) = &self.gpu {
|
||||
g.ok.store(false, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
fn gpu_failed(&self, what: &str, e: &Error) {
|
||||
if let Some(g) = &self.gpu
|
||||
&& g.ok.swap(false, Ordering::Relaxed)
|
||||
{
|
||||
log::warn!("SAM 2 {what} failed on the GPU ({e}); the CPU takes over");
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode a normalised `[1, 3, 1024, 1024]` frame on the GPU if it is working, else on the CPU.
|
||||
pub fn encode(&self, pixels: &Tensor) -> Result<Encoded> {
|
||||
if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) {
|
||||
match g.encoder.encode(pixels) {
|
||||
Ok(f) => {
|
||||
let none = || Tensor::zeros((1, 0, 0, 0), DType::F32, &candle_core::Device::Cpu);
|
||||
let pos = crate::neck::sine_position(128, GRID, GRID)?;
|
||||
let features = ImageFeatures { features: vec![none()?, none()?, f.low], positions: vec![none()?, none()?, pos] };
|
||||
return Ok(Encoded { features, high: f.high });
|
||||
}
|
||||
Err(e) => self.gpu_failed("encoding", &e),
|
||||
}
|
||||
}
|
||||
self.image.encode(pixels)
|
||||
}
|
||||
|
||||
/// The memory of a frame (its features and the scaled mask), on the GPU if it is working, else on the CPU.
|
||||
fn memorize(&self, features: &Tensor, mask: &Tensor, present: bool) -> Result<Memory> {
|
||||
if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) {
|
||||
match g.memory_encoder.encode(features, mask).and_then(|t| self.memory_encoder.finish(&t, present)) {
|
||||
Ok(m) => return Ok(m),
|
||||
Err(e) => self.gpu_failed("memory encoding", &e),
|
||||
}
|
||||
}
|
||||
self.memory_encoder.memory(features, mask, present)
|
||||
}
|
||||
|
||||
/// Memory attention on the GPU if it is working, else on the CPU.
|
||||
fn attend(&self, current: &Tensor, current_pos: &Tensor, memory: &Tensor, memory_pos: &Tensor, pointers: usize) -> Result<Tensor> {
|
||||
if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) {
|
||||
match g.attention.forward(current, current_pos, memory, memory_pos, pointers) {
|
||||
Ok(t) => return Ok(t),
|
||||
Err(e) => self.gpu_failed("memory attention", &e),
|
||||
}
|
||||
}
|
||||
self.memory_attention.forward(current, current_pos, memory, memory_pos, pointers)
|
||||
}
|
||||
|
||||
pub fn open(path: &Path) -> Result<VideoModel> {
|
||||
VideoModel::load(&Weights::open(path)?)
|
||||
}
|
||||
@@ -221,7 +350,7 @@ impl<'a> VideoTracker<'a> {
|
||||
return Ok(f.clone());
|
||||
}
|
||||
let pixels = self.frames.frame(frame)?;
|
||||
let enc = self.model.image.encode(&pixels)?;
|
||||
let enc = self.model.encode(&pixels)?;
|
||||
let (low, pos) = match (enc.features.features.get(2), enc.features.positions.get(2)) {
|
||||
(Some(l), Some(p)) => (l.clone(), p.clone()),
|
||||
_ => return Err(Error::Compute("vision encoder returned too few levels".into())),
|
||||
@@ -295,7 +424,17 @@ impl<'a> VideoTracker<'a> {
|
||||
|
||||
/// Run the model on one frame for one object and store the result.
|
||||
fn step(&mut self, id: ObjectId, frame: usize, prompt: Option<&Prompt>, reverse: bool, conditioning: bool) -> Result<FrameMask> {
|
||||
// EFFECTCRAFT_SAM2_PROFILE=1 prints where a frame's time goes (stderr).
|
||||
let profile = std::env::var_os("EFFECTCRAFT_SAM2_PROFILE").is_some();
|
||||
let mut lap = std::time::Instant::now();
|
||||
let mut mark = |what: &str| {
|
||||
if profile {
|
||||
eprintln!("sam2 frame {frame}: {what} {:.0} ms", lap.elapsed().as_secs_f64() * 1e3);
|
||||
}
|
||||
lap = std::time::Instant::now();
|
||||
};
|
||||
let feats = self.features(frame)?;
|
||||
mark("encode");
|
||||
let model = self.model.clone();
|
||||
let tokens = |t: &Tensor| -> Result<Tensor> { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) };
|
||||
let conditioned = if conditioning {
|
||||
@@ -303,9 +442,10 @@ impl<'a> VideoTracker<'a> {
|
||||
} else {
|
||||
let obj = self.objects.get(&id).ok_or_else(|| Error::Input("unknown object".into()))?;
|
||||
let (mem, mem_pos, n_ptr) = self.memory_bank(obj, frame, reverse)?;
|
||||
let out = model.memory_attention.forward(&tokens(&feats.low)?, &tokens(&feats.pos)?, &mem, &mem_pos, n_ptr)?;
|
||||
let out = model.attend(&tokens(&feats.low)?, &tokens(&feats.pos)?, &mem, &mem_pos, n_ptr)?;
|
||||
out.t()?.contiguous()?.reshape((1, 256, GRID, GRID))?
|
||||
};
|
||||
mark("memory attention");
|
||||
let num_points = prompt.map_or(0, |p| p.points.len() + 2 * usize::from(p.bbox.is_some()));
|
||||
let multimask = num_points <= 1;
|
||||
let sparse = match prompt {
|
||||
@@ -315,6 +455,7 @@ impl<'a> VideoTracker<'a> {
|
||||
let dense = model.image.prompt.dense_empty()?;
|
||||
let out = model.image.decode_embedded(&conditioned, &feats.high, &sparse, &dense, multimask)?;
|
||||
|
||||
mark("decoder");
|
||||
let score = scalar(&out.object_score)?;
|
||||
let present = score > 0.0;
|
||||
let ious = out.iou.flatten_all()?.to_vec1::<f32>()?;
|
||||
@@ -327,7 +468,9 @@ impl<'a> VideoTracker<'a> {
|
||||
let high = low.upsample_bilinear2d(INPUT_SIZE, INPUT_SIZE, false)?;
|
||||
let for_memory = if prompt.is_some() { high.gt(0f32)?.to_dtype(DType::F32)? } else { sigmoid(&high)? };
|
||||
let for_memory = ((for_memory * 20.0)? - 10.0)?;
|
||||
let memory = model.memory_encoder.memory(&feats.low, &for_memory, present)?;
|
||||
mark("upsample and mask");
|
||||
let memory = model.memorize(&feats.low, &for_memory, present)?;
|
||||
mark("memory encoder");
|
||||
|
||||
let iou = ious.get(best).copied().unwrap_or(0.0);
|
||||
let result = FrameMask { frame, object: id, logits: low.flatten_all()?.to_vec1::<f32>()?, iou, object_score: score, occluded: !present };
|
||||
@@ -414,7 +557,7 @@ impl<'a> VideoTracker<'a> {
|
||||
let decoder_present = scalar(&out.object_score)? > 0.0;
|
||||
let pointer = if decoder_present { model.pointer_proj.forward(&token)? } else { model.no_object_pointer.clone() };
|
||||
let pointer = if present { pointer } else { model.no_object_pointer.clone() };
|
||||
let memory = model.memory_encoder.memory(&feats.low, &high, present)?;
|
||||
let memory = model.memorize(&feats.low, &high, present)?;
|
||||
let score = if present { 10.0 } else { -10.0 };
|
||||
let result = FrameMask { frame, object: id, logits: low, iou: 1.0, object_score: score, occluded: !present };
|
||||
let state = self.objects.entry(id).or_default();
|
||||
|
||||
@@ -223,12 +223,17 @@ fn iou(a: &[f32], b: &[f32]) -> f64 {
|
||||
const VIDEO_LOGIT_TOL: f32 = 0.02;
|
||||
const VIDEO_MIN_IOU: f64 = 0.995;
|
||||
|
||||
#[test]
|
||||
fn video_tracking_matches_sam2_video_model() {
|
||||
fn check_video_tracking(backend: effectcraft_sam2::session::Backend) {
|
||||
use effectcraft_sam2::prompt::Prompt;
|
||||
use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker};
|
||||
let Some((w, r)) = load_video() else { return };
|
||||
let model = std::sync::Arc::new(VideoModel::load(&w).unwrap());
|
||||
let model = std::sync::Arc::new(match VideoModel::load_with(&w, backend) {
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
eprintln!("skipped: {e}");
|
||||
return;
|
||||
}
|
||||
});
|
||||
let clip = r.get("clip.pixel_values").unwrap().clone();
|
||||
let mut tracker = VideoTracker::new(model, Box::new(Clip(clip))).unwrap();
|
||||
let obj = ObjectId(1);
|
||||
@@ -261,6 +266,19 @@ fn video_tracking_matches_sam2_video_model() {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn video_tracking_matches_sam2_video_model() {
|
||||
check_video_tracking(effectcraft_sam2::session::Backend::Cpu);
|
||||
}
|
||||
|
||||
/// The same clip on the GPU: every frame's mask against `Sam2VideoModel`, forwards and backwards. The GPU path
|
||||
/// is within 1e-5 of PyTorch stage by stage; through the clip the bfloat16 memories can round the other way
|
||||
/// on a few values, so the bound on logits is the video one, and the masks must agree.
|
||||
#[test]
|
||||
fn video_tracking_on_the_gpu_matches_sam2_video_model() {
|
||||
check_video_tracking(effectcraft_sam2::session::Backend::Gpu);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracker_rejects_hostile_input_and_honours_cancel() {
|
||||
use effectcraft_sam2::Error;
|
||||
@@ -603,8 +621,8 @@ fn gpu_encoder_matches_the_cpu_path_and_pytorch() {
|
||||
use effectcraft_sam2::gpu::{Gpu, encoder::GpuEncoder};
|
||||
use effectcraft_sam2::model::ImageModel;
|
||||
let Some((w, r)) = load() else { return };
|
||||
let gpu = match Gpu::new() {
|
||||
Ok(g) => std::sync::Arc::new(g),
|
||||
let gpu = match Gpu::shared() {
|
||||
Ok(g) => g,
|
||||
Err(e) => {
|
||||
eprintln!("skipped: {e}");
|
||||
return;
|
||||
@@ -632,3 +650,105 @@ fn gpu_encoder_matches_the_cpu_path_and_pytorch() {
|
||||
eprintln!("skip features vs CPU: {e0:e} / {e1:e}");
|
||||
assert!(e0 < GPU_TOL && e1 < GPU_TOL, "skip feature errors {e0} / {e1}");
|
||||
}
|
||||
|
||||
/// The GPU memory attention against PyTorch (the fixture's frame: one memory frame and four pointer tokens).
|
||||
/// Measured on the RX 7900 GRE, 2026-10-15; the CPU path is at 2.6e-5.
|
||||
const GPU_ATTN_TOL: f32 = 1e-3;
|
||||
|
||||
#[test]
|
||||
fn gpu_memory_attention_matches_pytorch() {
|
||||
use effectcraft_sam2::gpu::{Gpu, memattn::GpuMemoryAttention};
|
||||
let Some((w, r)) = load_video() else { return };
|
||||
let gpu = match Gpu::shared() {
|
||||
Ok(g) => g,
|
||||
Err(e) => {
|
||||
eprintln!("skipped: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let att = GpuMemoryAttention::load(gpu, &w).unwrap();
|
||||
let get = |k: &str| r.get(k).unwrap().squeeze(1).unwrap();
|
||||
let pointers = r.get("mem.attn.num_pointers").unwrap().to_vec1::<f32>().unwrap()[0] as usize;
|
||||
let args = (
|
||||
get("mem.attn.current_vision_features"),
|
||||
get("mem.attn.current_vision_position_embeddings"),
|
||||
get("mem.attn.memory"),
|
||||
get("mem.attn.memory_posision_embeddings"),
|
||||
);
|
||||
let out = att.forward(&args.0, &args.1, &args.2, &args.3, pointers).unwrap();
|
||||
let want = r.get("mem.attn.out").unwrap().squeeze(0).unwrap().squeeze(0).unwrap();
|
||||
let e = max_abs(&out, &want);
|
||||
eprintln!("GPU memory attention vs PyTorch: max abs error {e:e}");
|
||||
assert!(e < GPU_ATTN_TOL, "GPU memory attention error {e}");
|
||||
// a full bank (7 memory frames and 16 pointers of 4 tokens): the same maths on a bigger problem, against the CPU path
|
||||
let cpu = effectcraft_sam2::memattn::MemoryAttention::load(&w).unwrap();
|
||||
let m = 7 * 4096 + 64;
|
||||
let mem = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu).unwrap();
|
||||
let mem_pos = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu).unwrap();
|
||||
let t = std::time::Instant::now();
|
||||
let big = att.forward(&args.0, &args.1, &mem, &mem_pos, 64).unwrap();
|
||||
eprintln!("GPU memory attention, full bank: {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
||||
let want = cpu.forward(&args.0, &args.1, &mem, &mem_pos, 64).unwrap();
|
||||
let e = max_abs(&big, &want);
|
||||
eprintln!("full bank vs CPU: max abs error {e:e}");
|
||||
assert!(e < GPU_ATTN_TOL, "full-bank error {e}");
|
||||
// bad shapes are errors
|
||||
assert!(att.forward(&args.0, &args.1, &mem, &mem_pos, m + 1).is_err());
|
||||
assert!(att.forward(&args.0.narrow(0, 0, 10).unwrap(), &args.1, &mem, &mem_pos, 64).is_err());
|
||||
}
|
||||
|
||||
/// The GPU can be switched off in the middle of a clip: the CPU takes over, the device name changes, and
|
||||
/// the masks still match `Sam2VideoModel` (the memory banks built on the GPU are read by the CPU path).
|
||||
#[test]
|
||||
fn switching_the_gpu_off_mid_clip_keeps_the_tracking_right() {
|
||||
use effectcraft_sam2::prompt::Prompt;
|
||||
use effectcraft_sam2::session::{Backend, Cancel, Direction, ObjectId, VideoModel, VideoTracker};
|
||||
let Some((w, r)) = load_video() else { return };
|
||||
let Ok(model) = VideoModel::load_with(&w, Backend::Gpu) else {
|
||||
eprintln!("skipped: no GPU");
|
||||
return;
|
||||
};
|
||||
let model = std::sync::Arc::new(model);
|
||||
assert_ne!(model.device_name(), "CPU");
|
||||
let mut t = VideoTracker::new(model.clone(), Box::new(Clip(r.get("clip.pixel_values").unwrap().clone()))).unwrap();
|
||||
t.prompt(2, ObjectId(1), &Prompt { points: vec![([416.0, 512.0], true)], bbox: None }).unwrap();
|
||||
let mut seen = vec![];
|
||||
t.propagate(2, Direction::Forward, Some(1), &Cancel::new(), &mut |m| seen.push(m)).unwrap();
|
||||
model.disable_gpu();
|
||||
assert_eq!(model.device_name(), "CPU");
|
||||
t.propagate(2, Direction::Forward, None, &Cancel::new(), &mut |m| seen.push(m)).unwrap();
|
||||
let last = seen.last().unwrap();
|
||||
assert_eq!(last.frame, 6);
|
||||
let want = r.get("video.fwd.6.pred_masks").unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap();
|
||||
let j = iou(&last.logits, &want);
|
||||
assert!(j >= VIDEO_MIN_IOU, "frame 6 after the switch: IoU {j}");
|
||||
}
|
||||
|
||||
/// The GPU memory encoder against the CPU one on the fixture's frame (a mask, the frame's features). The CPU
|
||||
/// path matches PyTorch to 3.5e-6; the GPU is held to the same output.
|
||||
#[test]
|
||||
fn gpu_memory_encoder_matches_the_cpu_one() {
|
||||
use effectcraft_sam2::gpu::{Gpu, memenc::GpuMemoryEncoder};
|
||||
let Some((w, r)) = load_video() else { return };
|
||||
let Ok(gpu) = Gpu::shared() else {
|
||||
eprintln!("skipped: no GPU");
|
||||
return;
|
||||
};
|
||||
let cpu = effectcraft_sam2::memory::MemoryEncoder::load(&w).unwrap();
|
||||
let enc = GpuMemoryEncoder::load(gpu, &w).unwrap();
|
||||
let (features, mask) = (r.get("mem.enc.vision_features").unwrap(), r.get("mem.enc.masks").unwrap());
|
||||
let t = std::time::Instant::now();
|
||||
let got = enc.encode(features, mask).unwrap();
|
||||
eprintln!("GPU memory encoder: {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
||||
let want = r.get("mem.enc.out_features").unwrap().flatten_from(2).unwrap().squeeze(0).unwrap().t().unwrap().contiguous().unwrap();
|
||||
let e = max_abs(&got, &want);
|
||||
eprintln!("GPU memory encoder vs PyTorch: max abs error {e:e}");
|
||||
assert!(e < 1e-3, "memory encoder error {e}");
|
||||
// the whole memory (with the occlusion embedding and bfloat16 rounding) equals the CPU one
|
||||
let (a, b) = (cpu.finish(&got, false).unwrap(), cpu.memory(features, mask, false).unwrap());
|
||||
let e = max_abs(&a.features, &b.features);
|
||||
assert!(e < 0.05, "memory features differ by {e} (bfloat16 steps are 2^-8 of the value)");
|
||||
assert_eq!(max_abs(&a.pos, &b.pos), 0.0);
|
||||
// wrong shapes are errors
|
||||
assert!(enc.encode(&features.narrow(1, 0, 8).unwrap(), mask).is_err());
|
||||
}
|
||||
|
||||
@@ -49,6 +49,10 @@ pub trait MaskModel: Send + Sync {
|
||||
fn tracker(&self) -> Option<&dyn ObjectTracker> {
|
||||
None
|
||||
}
|
||||
/// What runs the model: the GPU's name or "CPU" (`None`: it does not say).
|
||||
fn device_name(&self) -> Option<String> {
|
||||
None
|
||||
}
|
||||
/// Seconds since the model was last used (`None`: it does not track use, so it is never unloaded
|
||||
/// for being idle). Large models free their memory after a quiet spell and load again on use.
|
||||
fn idle_secs(&self) -> Option<f64> {
|
||||
|
||||
+21
-13
@@ -64,20 +64,28 @@ the CPU.
|
||||
SAM 2 is prompted with at most three foreground and three background points per frame: it was trained on a handful
|
||||
of clicks, and measured confidence collapses with many (0.99 with 1–7 points, 0.25 with 16, 0.00 with 32).
|
||||
|
||||
## Speed (this PC: 16 threads, CPU only, release build)
|
||||
| Step | Time |
|
||||
|---|---|
|
||||
| Encode one 1024 × 1024 frame | about **2.1 s** (3.4 s before the speed work: a parallel GELU, cached position encodings, pooling without permutes, a parallel bias add) |
|
||||
| A click on an encoded frame | about 100 ms |
|
||||
| Memory attention on a tracked frame (full bank of 7 memory frames) | about **2.3 s** |
|
||||
| A tracked frame in all | about **4.6 s** (the first six have fewer memories) |
|
||||
## Speed (this PC: AMD RX 7900 GRE and a 16-thread CPU, release build)
|
||||
SAM 2 runs on the GPU when there is one (wgpu compute kernels, so AMD, NVIDIA and Intel alike; candle itself has no AMD
|
||||
backend) and on the CPU otherwise. The model loads on the GPU automatically; `EFFECTCRAFT_SAM2_DEVICE=cpu` (or `gpu`)
|
||||
overrides it. A GPU that fails (no adapter, out of memory, a lost device) makes the CPU take over for the rest of the
|
||||
session, and `roto.status` reports `tracker.device` (the GPU's name, or `CPU`).
|
||||
|
||||
A 300-frame clip therefore takes on the order of 23 minutes on this CPU. (An earlier version of this page said one encode
|
||||
per frame; the memory attention is as large as the encode.) What was tried and did not help: chunking the queries of the
|
||||
global and memory attention to keep scores in cache. Candle has no AMD/Vulkan backend, so a GPU path is future work. A spike (`--example gpu_matmul`) measured a plain tiled
|
||||
WGSL matmul on the RX 7900 GRE at 1,900–3,600 GFLOPS against 160–470 for candle's CPU matmul on the same shapes (6–13×),
|
||||
so hand-written wgpu kernels could bring a tracked frame to a second or less; `cargo run --release -p effectcraft-sam2 --example sam2_bench` and `--example op_bench` time the
|
||||
pieces.
|
||||
| Step | CPU | GPU |
|
||||
|---|---|---|
|
||||
| Encode one 1024 × 1024 frame | 2.05 s | **0.24 s** |
|
||||
| Memory attention (full bank of 7 memory frames) | 2.3 s | **0.29 s** |
|
||||
| Memory encoder | 0.24 s | 0.05 s |
|
||||
| Decoder (CPU on both) | 0.11 s | 0.11 s |
|
||||
| **A tracked frame in all** | **4.3 s** | **0.7 s** |
|
||||
| A click on an encoded frame | 0.10 s | 0.10 s |
|
||||
|
||||
A 300-frame clip takes about 21 minutes on the CPU and about 3.5 minutes on this GPU; the decoder, still on the CPU, is the
|
||||
next 16 % of a GPU frame. The CPU numbers came down from 3.4 s per encode at the start of the speed
|
||||
work (parallel GELU and bias add, cached position encodings, pooling without permutes); chunking the attention queries was
|
||||
tried and did not help. The GPU kernels (`crates/sam2/src/gpu`) are tested against CPU definitions and PyTorch: the
|
||||
encoder's 64 × 64 level matches PyTorch to 7e-6, the memory attention to 1.8e-5, and a clip tracked on the GPU matches
|
||||
`Sam2VideoModel` as the CPU does (worst logit error 0.0032). `cargo run --release -p effectcraft-sam2 --example sam2_bench`
|
||||
times the pieces on both devices (`EFFECTCRAFT_SAM2_PROFILE=1` prints a frame's breakdown).
|
||||
|
||||
## Benchmark
|
||||
`effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt]` runs synthetic clips
|
||||
|
||||
Reference in New Issue
Block a user