README.md

---
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`