--- license: cc-by-nc-4.0 library_name: pytorch pipeline_tag: mask-generation tags: - video-object-segmentation - sam2 - knowledge-distillation - onnx --- # sam2-lite **SAM 2.1 video object segmentation 2.13× faster on CPU**: SAM 2.1 Hiera-tiny with its image encoder distilled into a MobileNetV4 backbone and its memory attention distilled to attend to 3 memory frames instead of 7. Click an object in the first frame and it is tracked through the video. Code, training and evaluation: https://github.com/pablopigue/sam2-lite ## Results DAVIS 2017 val (30 videos, semi-supervised, J&F), per-frame latency of the full video pipeline (median, fp32, 13th Gen Intel(R) Core(TM) i7-13620H). | | SAM 2.1 Hiera-T (teacher) | **sam2-lite** @1024 | **sam2-lite-mobile** @576 | |---|---|---|---| | J&F | 89.1 | 84.4 | 77.1 | | CPU 6 threads (ms/frame) | 1956 | 920 | – | | CPU 2 threads (ms/frame) | 4949 | 2270 | 398 | | Model size, fp32 (MB) | 149 | 74 | 74 | sam2-lite-mobile is the same weights at a 576×576 input; it was measured in a different session from the other latency numbers (~10 % run-to-run variation). 2 CPU threads = a modest 2-core CPU. ### ONNX Runtime encoder The image encoder is also provided as ONNX (one file per resolution). PyTorch → ONNX Runtime, all measured in one later session (hence e.g. 892 instead of 920 ms for the same PyTorch model: ~5-10 % varies between sessions, so only same-session numbers are compared): | | Encoder (ms) | Full pipeline (ms) | |---|---|---| | 1024, CPU 6 threads | 224 → 142 (1.58×) | 892 → 830 (1.07×) | | 1024, CPU 2 threads | 434 → 295 (1.47×) | 2236 → 2075 (1.08×) | | 576, CPU 6 threads | 44 → 40 (1.11×) | 159 → 151 (1.06×) | | 576, CPU 2 threads | 96 → 75 (1.28×) | 392 → 373 (1.05×) | The encoder is faster but the pipeline gains little: the memory attention, which still runs in PyTorch, dominates each frame. Quality is unchanged (5 val videos at 576: J&F 82.35 with PyTorch, 82.35 with ONNX Runtime). ## Files - `model.safetensors` + `config.yaml`: the whole tracker (student encoder, fine-tuned memory attention, SAM 2.1-tiny memory encoder and mask decoder) and what is needed to rebuild it. - `student_encoder_r1024.onnx`, `student_encoder_r576.onnx`: the image encoder for ONNX Runtime (CPU). - `LICENSE`, `LICENSE-APACHE-2.0`, `NOTICE`: see [License](#license). ## Usage ```bash git clone https://github.com/pablopigue/sam2-lite && cd sam2-lite && uv sync ``` ```python from huggingface_hub import snapshot_download from sam2lite.export.bundle import load_bundle path = snapshot_download("Pablopigue/sam2-lite") predictor = load_bundle(path) # sam2-lite @1024 mobile = load_bundle( path, image_size=576, onnx=f"{path}/student_encoder_r576.onnx" ) # sam2-lite-mobile, ONNX encoder # Then use it like SAM 2's video predictor: init_state, add_new_points_or_box, propagate_in_video. ``` ## Training - **Image encoder:** MobileNetV4-Conv-Medium (timm, ImageNet-1k pretrained) + SAM 2's FPN neck, trained to reproduce the teacher encoder's multi-scale features (MSE per FPN level, FitNets-style), on ~3.8k unlabeled frames from 54 DAVIS 2017 train videos. The rest of SAM 2.1 is reused unchanged. - **Memory attention:** a copy of SAM 2.1's, fine-tuned to reproduce the original's output (7 memory frames) while attending to 3 (stride 4), on DAVIS 2017 train clips. - DAVIS 2017 val was used only for the final evaluation. ## Limitations - Accuracy drops by 4.7 J&F points at 1024; a few small or close-up objects explain most of the drop. - sam2-lite-mobile misses its own quality target (J&F ≥ 80) while meeting its latency target (≤ 600 ms with 2 threads). - Distilled on a small dataset (DAVIS 2017 train only); the experiments point at the data as the current bottleneck. - Latency was measured on one laptop CPU; numbers on other hardware will differ. ## License sam2-lite is an independent project; it is not affiliated with, endorsed by or sponsored by Meta. "SAM 2" refers to the original model by Meta FAIR, on which this work is based. Released under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/) for **non-commercial research use only** (`LICENSE`); `NOTICE` details every part: - It contains SAM 2.1 weights by Meta Platforms, Inc., licensed under the Apache License 2.0 (`LICENSE-APACHE-2.0`). **Modified:** the memory attention (fine-tuned to attend to 3 memory frames instead of 7) and the memory temporal encodings (3 of the 7 kept). The other SAM 2.1 weights (memory encoder, prompt encoder, mask decoder, object pointers) are unchanged and remain available under the Apache License 2.0. - The image encoder starts from the timm weights [`mobilenetv4_conv_medium.e500_r224_in1k`](https://huggingface.co/timm/mobilenetv4_conv_medium.e500_r224_in1k) (Apache 2.0), **modified** by distillation. They were pretrained on ImageNet-1k, whose terms allow only non-commercial research and educational use. - It was distilled on frames of [DAVIS 2017](https://davischallenge.org), licensed under CC BY-NC 4.0; about half of its sequences come from third-party sources (mostly YouTube) with their own terms. No DAVIS videos, frames or annotations are included. ## Citations ```bibtex @article{ravi2024sam2, title = {SAM 2: Segment Anything in Images and Videos}, author = {Ravi, Nikhila and others}, journal = {arXiv preprint arXiv:2408.00714}, year = {2024} } @article{pont2017davis, title = {The 2017 DAVIS Challenge on Video Object Segmentation}, author = {Pont-Tuset, Jordi and Perazzi, Federico and Caelles, Sergi and Arbel{\'a}ez, Pablo and Sorkine-Hornung, Alexander and Van Gool, Luc}, journal = {arXiv preprint arXiv:1704.00675}, year = {2017} } @inproceedings{perazzi2016davis, title = {A Benchmark Dataset and Evaluation Methodology for Video Object Segmentation}, author = {Perazzi, Federico and Pont-Tuset, Jordi and McWilliams, Brian and Van Gool, Luc and Gross, Markus and Sorkine-Hornung, Alexander}, booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, year = {2016} } @article{qin2024mobilenetv4, title = {MobileNetV4: Universal Models for the Mobile Ecosystem}, author = {Qin, Danfeng and others}, journal = {arXiv preprint arXiv:2404.10518}, year = {2024} } ``` MLflow runs behind every number: `b44efdd1`, `62e59f5e`, `ca473c99`, `a0b4ebcc`, `a384f235`, `b205e7be`, `af0d9753`, `7baca321`, `79a6f48c`, `005f6d98`, `7be93e7d`, `4e29d891`