mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-01 05:58:19 -06:00
Kristian MyramandClaude Sonnet 5.5
d6fcd4a62b
A tracked frame is the encode (2.4 s) plus memory attention (2.4 s with a full bank), about 5 s on 16 threads, not one encode as the docs said. Docs corrected; chunked attention tried and reverted. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
76 lines
3.5 KiB
Rust
76 lines
3.5 KiB
Rust
//! Timing of the SAM 2 image path on this machine: load, encode one frame, decode a click.
|
|
//!
|
|
//! `cargo run -p effectcraft-sam2 --example sam2_bench -- [MODEL_DIR]` (default `sam2-tiny`).
|
|
//! Prints a message and exits when the weights are not there.
|
|
|
|
use std::time::Instant;
|
|
|
|
use candle_core::{DType, Device, Tensor};
|
|
use effectcraft_sam2::model::ImageModel;
|
|
use effectcraft_sam2::prompt::Prompt;
|
|
|
|
fn main() {
|
|
if let Err(e) = run() {
|
|
eprintln!("sam2_bench: {e}");
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
|
|
fn run() -> Result<(), Box<dyn std::error::Error>> {
|
|
let dir = std::env::args().nth(1).unwrap_or_else(|| "sam2-tiny".to_string());
|
|
let file = std::path::Path::new(&dir).join("model.safetensors");
|
|
if !file.exists() {
|
|
println!("no weights at {} (download facebook/sam2.1-hiera-tiny first)", file.display());
|
|
return Ok(());
|
|
}
|
|
let t = Instant::now();
|
|
let model = ImageModel::open(&file)?;
|
|
println!("load: {:.2} s", t.elapsed().as_secs_f64());
|
|
let pixels = Tensor::randn(0f32, 1.0, (1, 3, 1024, 1024), &Device::Cpu)?.to_dtype(DType::F32)?;
|
|
for run in 0..3 {
|
|
let t = Instant::now();
|
|
let enc = model.encode(&pixels)?;
|
|
let encode = t.elapsed().as_secs_f64();
|
|
let prompt = Prompt { points: vec![([512.0, 512.0], true)], bbox: None };
|
|
let t = Instant::now();
|
|
let out = model.segment(&enc, &prompt, true)?;
|
|
println!("run {run}: encode {encode:.2} s, click {:.0} ms, masks {:?}", t.elapsed().as_secs_f64() * 1e3, out.masks.dims());
|
|
}
|
|
// Where the encode time goes: patch embedding, then each Hiera block.
|
|
let hiera = &model.vision.hiera;
|
|
let t = Instant::now();
|
|
let mut x = hiera.patch_embed(&pixels)?;
|
|
println!("patch embed: {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
|
for (i, b) in hiera.blocks.iter().enumerate() {
|
|
let t = Instant::now();
|
|
x = b.forward(&x)?;
|
|
println!("block {i}: {:.0} ms, out {:?}", t.elapsed().as_secs_f64() * 1e3, x.dims());
|
|
}
|
|
// The rest of the encode: the neck (1x1 convolutions, top-down sums) and its sine position encodings.
|
|
let stages = hiera.forward(&pixels)?;
|
|
let t = Instant::now();
|
|
let _ = model.vision.neck.forward(&stages)?;
|
|
println!("neck (convolutions + positions): {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
|
let t = Instant::now();
|
|
for (h, w) in [(32, 32), (64, 64), (128, 128), (256, 256)] {
|
|
let _ = effectcraft_sam2::neck::sine_position(128, h, w)?;
|
|
}
|
|
println!("sine positions for the four levels: {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
|
// What a tracked frame adds to the encode: memory attention over the memory bank (7 memory frames and
|
|
// 16 object pointers of 4 tokens), then the memory encoder.
|
|
let weights = effectcraft_sam2::weights::Weights::open(&file)?;
|
|
let attn = effectcraft_sam2::memattn::MemoryAttention::load(&weights)?;
|
|
let cur = Tensor::randn(0f32, 1.0, (4096, 256), &Device::Cpu)?;
|
|
let cur_pos = Tensor::randn(0f32, 1.0, (4096, 256), &Device::Cpu)?;
|
|
let m = 7 * 4096 + 64;
|
|
let mem = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu)?;
|
|
let mem_pos = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu)?;
|
|
for run in 0..2 {
|
|
let t = Instant::now();
|
|
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);
|
|
}
|
|
println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get()));
|
|
Ok(())
|
|
}
|