SAM 2 from Scratch: Promptable Video Segmentation in the Browser

01 Overview

The reimplementation covers every inference-time module: the Hiera hierarchical vision transformer with windowed and global attention and query pooling between stages, the feature pyramid neck, the prompt encoder with its random Fourier position encoding, the two-way transformer mask decoder with its IoU and occlusion heads, the memory encoder that fuses a predicted mask back into image features, the memory attention that conditions each new frame on the prompted frames, the six most recent frames and up to sixteen object pointers, with two-dimensional rotary position embeddings, and the propagation loop that decides which memories a frame attends to. Training on the 50,900-video SA-V dataset took Meta 256 A100 GPUs for 108 hours; that part is not reproduced, and this page does not pretend otherwise.

Correctness is defined numerically rather than by eye. A dump tool records every intermediate tensor the official package produces on three test images and a 48-frame clip, including the exact inputs and outputs of the memory attention at five points in the video. Each of my modules keeps the parameter names of its checkpoint slice, loads it strictly, and is tested against those dumps. The result on both released sizes, Hiera-T and Hiera-S, is a maximum absolute error of exactly zero on every module and every propagated frame, forward and backward in time. Two details that the paper and configuration files do not state had to be recovered from the numbers: masks of clicked frames are binarised before they enter memory, and memory features are stored in bfloat16, which changes later frames by more than the test tolerance if skipped.

To check the model against a public benchmark I ran the standard semi-supervised protocol on DAVIS 2017 val, where every object is given its first-frame mask and tracked to the end of the clip. Hiera-T scores 87.7 J&F and Hiera-S 88.0. Meta's appendix reports 89.4 and 89.6 for the same architectures, but for models trained with additional internal data that were never released. On three of the sequences re-run with the official package and the public SAM 2.1 checkpoints, Meta's code matches mine to within 0.002 points, so whatever the gap is, it is not in the implementation.

The demo below runs the full model in your browser with ONNX Runtime Web, preferring WebGPU and falling back to WebAssembly. It runs at 512 pixels rather than the paper's 1024: on DAVIS that costs 5.6 J&F (82.1 against 87.7) and makes every frame about four times cheaper, which is the trade a browser demo has to make. Six graphs are exported from my implementation rather than Meta's: the image encoder, the prompt-and-mask head, the memory encoder, the memory attention, and two small object-pointer graphs. The 27-million-parameter encoder is the download problem. Standard dynamic int8 quantisation dropped mask agreement with the fp32 model to 0.93 IoU and uses operators that have no WebGPU kernels, so the shipped encoder uses weight-only int8 with plain dequantisation nodes, keeps the patch embedding and first two blocks in fp32, and agrees with fp32 at 0.99 IoU in 32.5 MB. The video loop was first rewritten in plain onnxruntime and numpy and shown to match the PyTorch predictor exactly, then ported to TypeScript and tested in Node against that loop frame by frame before it touched the page.

What is still open, in the order I would attack it: tracking several objects at once in the browser, which the PyTorch side already does by batching memories; box and mask prompts in video mode, which the exported head already accepts; a 1024-pixel video mode, which needs an int8 encoder that keeps parity on frames where two mask candidates are nearly tied, since the current one flips between them; and a small fine-tuning study of the mask decoder on a narrow domain to see how far the frozen encoder carries.

05 What was built