Mask Generation
PyTorch
ONNX
Safetensors
sam2
video-object-segmentation
knowledge-distillation

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.

Usage

git clone https://github.com/pablopigue/sam2-lite && cd sam2-lite && uv sync
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 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 (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, 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

@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

Downloads last month
17
Safetensors
Model size
19.3M params
Tensor type
F32
Β·
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Papers for Pablopigue/sam2-lite