litert-community/DINOv2-ViT-S14-LiteRT

DINOv2 ViT-S/14 — dense features on LiteRT GPU

The self-supervised DINOv2 ViT-S/14 backbone running its full forward pass on the LiteRT CompiledModel GPU delegate (no CPU fallback). It emits the dense patch tokens; a top-3 PCA of those tokens mapped to RGB gives the classic "what the backbone sees" overlay — semantically similar patches (object parts vs background) share a color, so the object pops out with no labels or segmentation.

• Architecture: DINOv2 ViT-S/14 (vit_small_patch14_dinov2 in timm) — 12 blocks, dim 384, 6 heads, patch 14. Fixed 448×448 input → 32×32 = 1024 patch tokens.
• Weights: facebookresearch/dinov2 · Apache-2.0.
• Size: 45 MB (fp16).

Left: input. Right: top-3 PCA of the DINOv2 patch tokens (on-device fp16 output).

I/O

• Input: [1, 3, 448, 448] NCHW, RGB, ImageNet-normalized.
• Output: [1, 1024, 384] patch tokens (32×32 grid, cls token dropped).

GPU conversion

Fully GPU-resident on a Pixel 8a (864/864 nodes, 1 partition, ~8 ms) via the proven ViT re-authorings:

• 4D attention: the fused-qkv attention is split into q/k/v and reshaped to [1, heads, N, d] (≤4D) with a manual softmax(qkᵀ/√d)·v; the delegate rejects the native 5D head-split reshape.
• SafeLayerNorm: the deviation is scaled by 1/64 before squaring so the per-token sum of squares stays in fp16 range on DINOv2's massive activations, then rescaled — algebraically identical to the plain variance.
• LayerScale (ls1/ls2) baked into the following projection weights.
• tanh-GELU (0.5x(1+tanh(0.79788(x+0.044715x³)))) — near-exact and delegate-friendly; the sigmoid-GELU approximation drifts to feature corr 0.968 over 12 blocks, tanh → 0.99999.
• The pos_embed is baked at a fixed 448 grid at model creation, so there is no runtime interpolation (no GATHER_ND).

Device fp16 patch features vs desktop fp32: corr 0.996. Re-authored torch vs stock timm: corr 0.999992.

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

val model = CompiledModel.create(context.assets, "dinov2_s_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()

inputs[0].writeFloat(imageNchw)          // [1,3,448,448] ImageNet-normalized
model.run(inputs, outputs)
val tokens = outputs[0].readFloat()      // [1024*384] patch tokens -> PCA host-side

Python (LiteRT CompiledModel API)

import numpy as np
from ai_edge_litert.compiled_model import CompiledModel

model = CompiledModel.from_file("dinov2_s_fp16.tflite")
inputs = model.create_input_buffers(0)
outputs = model.create_output_buffers(0)
inputs[0].write(np.ascontiguousarray(image, np.float32))   # [1,3,448,448]
model.run_by_index(0, inputs, outputs)
tokens = outputs[0].read(1024 * 384, np.float32).reshape(1024, 384)

x = tokens - tokens.mean(0)
_, _, vt = np.linalg.svd(x, full_matrices=False)
rgb = x @ vt[:3].T                        # [1024,3] -> normalize -> 32x32 RGB

License

Apache-2.0 (DINOv2 / Meta). Converted with litert-torch.