File size: 4,835 Bytes
82972af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | ---
license: apache-2.0
library_name: litert
pipeline_tag: image-feature-extraction
tags:
- litert
- tflite
- sam2
- segment-anything
- image-encoder
- on-device
- gpu
base_model: facebook/sam2.1-hiera-tiny
---
# SAM 2.1 (Hiera-Tiny) image encoder — LiteRT GPU
On-device **LiteRT / TFLite** conversion of the **image encoder** of
[**SAM 2.1 Hiera-Tiny**](https://huggingface.co/facebook/sam2.1-hiera-tiny) (Meta, Apache-2.0),
running **fully on the mobile GPU** via the LiteRT `CompiledModel` API (ML Drift / `LITERT_CL` delegate).
The whole graph is GPU-resident — no CPU/XNNPACK fallback ops.
This is the heavy backbone of the Segment Anything 2 image path: it turns an RGB image into the
multi-scale feature pyramid that a (small) prompt-encoder + mask-decoder then query per click/box.
| | |
|---|---|
| Task | Image encoder for promptable segmentation (SAM 2 image path) |
| Backbone | Hiera-Tiny (hierarchical ViT, window + global attention) + FPN neck |
| Input | `[1, 3, 1024, 1024]` NCHW float32, ImageNet-normalized |
| Outputs | 3 FPN feature maps: `[1,256,256,256]`, `[1,256,128,128]`, `[1,256,64,64]` |
| Precision / size | FP16, **80 MB** |
| Device | Pixel 8a, LiteRT GPU (`Accelerator.GPU`), **~7 ms / image** |
| Residency | **`Replacing 862 out of 862 node(s) with delegate (LITERT_CL)`** (full, single partition) |
## Preprocessing (must match)
```
resize to 1024x1024 (bilinear) -> x/255 -> (x - mean) / std
mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225] # ImageNet, RGB, NCHW
```
## GPU-clean conversion (what was re-authored)
Converted with `litert-torch`. SAM 2's Hiera encoder is not GPU-clean out of the box; these exact,
weights-faithful rewrites were applied (model-side only — **no converter patch**):
1. **`window_partition` / `window_unpartition`**: the 6-D `view`+`permute` window reshape rejected by the
GPU delegate (>4-D) is re-expressed as a sequence of **≤4-D** `reshape`/`transpose` ops (numerically
exact, verified vs the original).
2. **`Sam2MultiScaleAttention`**: the 5-D fused-QKV reshape is decomposed into separate q/k/v, and
attention runs as a **3-D batched SDPA** (`[B*heads, N, d]`). A 4-D SDPA makes the delegate emit a
`[C,C]->[nW,ws,C,C]` `BROADCAST_TO` on every windowed block; the 3-D form removes all 9.
3. **Windowed positional embedding**: the bicubic-interpolate + tile of the constant `pos_embed` is
**baked to a buffer** (add only) — removes a runtime interpolate of a constant.
4. **Neck**: the (constant, shape-only) sine FPN position encodings are dropped from the graph (compute
them host-side) — removes the remaining `BROADCAST_TO` ops.
5. **Overflow-safe LayerNorm** (scale-before-square) as an fp16 safety margin for the deep stages.
Net: `banned ops = NONE`, `>4-D tensors = 0`, full GPU residency.
## Fidelity (honest)
Eager re-authoring is **numerically exact** (`cos = 1.000`, `mae = 0`). On-device GPU output vs the
CPU reference, per FPN level:
| Output | cosine |
|---|---|
| FPN-0 `256x256` (high-res, drives mask detail) | **0.99998** |
| FPN-1 `128x128` | **0.99994** |
| FPN-2 `64x64` (coarse image embedding) | **0.99253** |
The deepest 64×64 feature drifts slightly on the GPU. This is **not** LayerNorm overflow
(scale-before-square LayerNorm doesn't change it, and the CPU fp16 model matches PyTorch fp32 at
corr 0.999999) — it is the mobile GPU computing the deep-stage global attention (64×64 = 4096 tokens)
in true fp16, where the CPU path upcasts to fp32. The high-resolution features that carry mask
boundaries are near-exact, so mask quality is preserved in practice.
## Usage (Android / LiteRT CompiledModel)
```kotlin
val model = CompiledModel.create(context.assets, "sam2_tiny_image_encoder_fp16.tflite",
CompiledModel.Options(Accelerator.GPU), null)
// input: [1,3,1024,1024] NCHW, ImageNet-normalized
// outputs: 3 FPN feature maps -> feed to the SAM 2 prompt encoder + mask decoder
```
## Training data & PII
SAM 2 was trained by Meta on **SA-1B** (licensed photos) and **SA-V** (licensed videos) with
model-in-the-loop mask annotation. No new training was performed for this conversion — it is a
weights-faithful format change of the public `facebook/sam2.1-hiera-tiny` checkpoint. Because the
source data is real-world imagery, it may incidentally contain people, faces, vehicles, signage and
other PII; no PII was deliberately collected and this conversion adds none. Apply your own content/PII
filtering as appropriate. See the [SAM 2 release](https://github.com/facebookresearch/sam2) and
[paper](https://arxiv.org/abs/2408.00714) for full dataset details.
## License
Apache-2.0, inherited from the upstream [SAM 2.1](https://huggingface.co/facebook/sam2.1-hiera-tiny).
This is a format conversion; all credit to the original authors (Meta AI).
|