mlboydaisuke/EdgeTAM-ExecuTorch

EdgeTAM — ExecuTorch XNNPACK (encoder + decoder)

Promptable segmentation in two .pte files: run the encoder once per image, the decoder once per click.

• edgetam_encoder_xnnpack_fp32.pte (19.7 MB) — image (1,3,1024,1024) → image_embed (1,256,64,64), feat_s0 (1,32,256,256), feat_s1 (1,64,128,128)
• edgetam_decoder_xnnpack_fp32.pte (24.7 MB) — (image_embed, feat_s0, feat_s1, points (1,1,N,2) fp32 pixel coords in 1024-space, labels (1,1,N) int64 1=fg/0=bg) → mask logits (1,1,3,256,256), iou scores (1,1,3)
• edgetam_decoder_xnnpack_fp16.pte (12.6 MB) — the same decoder at half the size, corr 1.000000 against fp32 eager. It takes and returns fp32 tensors, so pairing it with the fp32 encoder needs no app changes.

The encoder ships in fp32 only, and that is not an omission. Its backbone is RepViT, which is convolutional, and XNNPACK serializes convolution weights as fp32 whatever dtype the graph carries, so fp16 came out at 19.8 MB (100.5%) and buys nothing.

Both int8 recipes were measured. Dynamic int8 annotates aten.linear and this graph has none — all 50 nn.Linear in EdgeTAM sit in the mask decoder, not on the image-embedding path — so it leaves the encoder at the same 19.7 MB. Static int8 does reach convolution weights, and on paper it is the build you would want: 5.4 MB and 24.7 ms against fp32's 19.7 MB and 31.0 ms on the same machine. It is the masks that stop it. Feeding its embeddings to the same fp32 decoder, over 37 clicks the fp32 arm answered, the worst decoded mask comes back at 0.0185 IoU against the fp32 encoder's — the median holds at 0.887, but a build that loses the mask outright on some clicks is not one to publish. Measured with a global per-channel static configuration calibrated on ten photographs; a narrower annotation might do better and has not been tried.

At 19.7 MB the fp32 encoder is already smaller than SAM 2.1 hiera-tiny's fp16 encoder (55.6 MB).

EdgeTAM is Meta's on-device SAM 2 (CVPR 2025). Its encoder is 5.5× smaller than SAM 2.1 hiera-tiny's (19.7 MB vs 109.2 MB) for the same output contract, so an app written against SAM2.1-hiera-tiny-ExecuTorch swaps the two files and changes nothing else.

• Source: facebook/EdgeTAM, loaded from the transformers-format mirror yonigozlan/EdgeTAM-hf (facebook/EdgeTAM publishes only the original edgetam.pt)
• License: Apache-2.0
• Preprocess: RGB/255, ImageNet norm (mean .485/.456/.406, std .229/.224/.225), resize 1024×1024
• Postprocess: take argmax(iou) of the 3 mask logits, threshold at > 0, upsample 4× (256→1024) back to image space. The prompt encoder is inside the decoder — pass raw click coordinates, no separate point-encoding code needed.

Verification (Mac arm64, executorch 1.4.0, torch 2.13.0)

Every output of both graphs matches torch fp32 eager at corr 1.000000, and the two wrappers compose back to EdgeTamModel.forward exactly (max_abs_diff 0.000e+00).

graph	output	shape	max_abs_diff	corr
encoder	image_embed	[1, 256, 64, 64]	0.000e+00	1.000000
encoder	feat_s0	[1, 32, 256, 256]	0.000e+00	1.000000
encoder	feat_s1	[1, 64, 128, 128]	0.000e+00	1.000000
decoder	mask logits	[1, 1, 3, 256, 256]	0.000e+00	1.000000
decoder	iou	[1, 1, 3]	0.000e+00	1.000000

Median over 10 runs, Mac arm64 single process — a relative reference, not a device number: encoder 32.3 ms (torch eager 103.5 ms), decoder 23.7 ms (eager 13.9 ms).

XNNPACK delegate coverage: encoder 99.8% (one upsample_nearest2d on the portable kernels), decoder 66.5% (the prompt encoder's expand/where bookkeeping stays on portable; every convolution and matmul is delegated).

Conversion

torch.export → to_edge_transform_and_lower(XnnpackPartitioner) → .pte (conversion script: executorch-models)

Two details matter for this split, both shared with the SAM 2.1 conversion. Encoder outputs are forced .contiguous() — transformers hands back channels_last tensors, and that layout at a .pte boundary makes the delegate's runtime shape propagation read physical strides as logical dims. Identity repeat_interleave(1, dim) calls in the decoder are dropped, since their lowered form mis-sizes on a single-point export.

The GPU-specific rewrites in the LiteRT build of this model (splitting the squeeze-excite mean, replacing ConvTranspose2d) are ML Drift workarounds and are not needed here — XNNPACK runs the stock graph.