rfdetr-coreai

Convert RF-DETR (Roboflow) object detectors to Core AI .aimodel assets for efficient inference on macOS and iOS.

This first milestone converts RF-DETR-Nano (384×384, pretrained COCO, 90 classes) and verifies the Core AI model numerically against the PyTorch source. The pipeline is parameterized to extend to the other variants (Small / Medium / Base / Large).

How it works

Follows Apple's coreai-models recipe pattern (models/yolo/export.py):

RFDETRNano() ─► deploy nn.Module (.export()) ─► [coreai-opt compression] ─► torch.export.export
            ─► run_decompositions ─► TorchConverter.to_coreai ─► optimize ─► save_asset(.aimodel)

The optional compression step (quantization / palettization) is selected per optimization profile — see Optimization profiles. The asset is then optionally AOT-compiled for the target platform with coreai-build — see Compile.

• The deploy model takes one image tensor (1, 3, 384, 384) (ImageNet-normalised, NCHW) and returns dets (1, 300, 4) (normalised cxcywh boxes) and labels (1, 300, 91) (class logits incl. a no-object slot). It is NMS-free; decoding is per-class sigmoid → threshold → cxcywh→xyxy (see decode.py).
• The logit column index equals the COCO category id (e.g. column 17 → cat).

The two export-compatibility shims

RF-DETR's deploy model exports fine through the legacy torch.onnx.export, but the stricter torch.export.export (required by coreai-torch) needs two tweaks, both localized to a single context manager in export_compat.py:

1. Data-dependent assert — MSDeformAttn.forward calls torch._assert(sum(H_l·W_l) == flat_len). The left side is tensor-derived, so torch.export treats it as an unbacked symbol and raises GuardOnDataDependentSymNode. It's pure input validation, so torch._assert is no-op'd during export (it's the only one in the model path).
2. F.grid_sample → aten.grid_sampler_2d — RF-DETR's deformable attention samples with F.grid_sample on CPU/CUDA (only its hand-written gather fallback runs on MPS). coreai-torch has no lowering for aten.grid_sampler_2d, so F.grid_sample is swapped for the same gather-based bilinear implementation. It's numerically identical (0.0 max-abs diff on the deploy outputs).

Setup

Requires uv and Python 3.12.

uv sync

This installs rfdetr, coreai-torch, coreai-core, and friends into .venv.

Dependency note (no quirk): coreai-torch==0.4.0 pins torch>=2.8,<=2.11 and rfdetr>=1.8 pins torch>=2.2; uv resolves both to torch==2.11.0 with no conflict, so the planned pip fallback was not needed. If a future version bump reintroduces a conflict, install with pip forcing coreai-torch's torch version and document it here.

Usage

Platform targeting (on-device authoring)

--platform tailors the asset to the target compute unit and sets a sensible default profile:

Platform	Compute unit	Default profile	Shapes	Goal
ios	Neural Engine	small (8-bit palett)	static, batch 1	energy / battery
macos	GPU	balanced (fp16)	fixed batch (--batch)	throughput

The preferred compute unit is recorded in asset metadata and applied at load time via SpecializationOptions.from_preferred_compute_unit_kind(...) — there's no build-time "prefer ANE" flag, so the consumer follows this contract (see verify.py). verify loads each asset with its platform's preferred unit automatically.

Dynamic batch: RF-DETR's deploy graph bakes the batch size at trace time (the deformable decoder's tensor-derived value.split(...)), and CoreAI's enumerated-shapes path crashes on it — so true 1..N batch isn't reachable without decoder surgery. Each asset is traced at one fixed batch (--batch, default 1); raise it for macOS/GPU throughput.

Perf caveat: the macOS-26 Python runtime is a ~1.3 s/inference reference path, so this repo verifies authoring correctness (detection agreement), not the ANE/GPU speed/energy payoff — confirm that on macOS 27 hardware. Deeper ANE authoring (BC1S / Conv2d rewrite of the ViT) is deferred for the same reason (see Roadmap).

Optimization profiles

--profile overrides the accuracy↔size tradeoff independently of platform (profiles.py):

Profile	dtype + compression	Asset size	Box drift	When to use
max-accuracy	fp32, none	~103 MB	0.00 px	reference / debugging
balanced (macOS default)	fp16, none	~52 MB	0.70 px	best all-round quality
small (iOS default)	8-bit palettized	~27 MB	0.75 px	near-lossless, smaller
tiny	6-bit palettized	~22 MB	1.02 px	size-critical on device

Convert (writes to exports/rfdetr-<variant>-<platform>[-<profile>][-b<N>].aimodel):

uv run rfdetr-convert --variant nano --platform ios                  # -> rfdetr-nano-ios-small.aimodel (~27 MB, ANE)
uv run rfdetr-convert --variant nano --platform macos                # -> rfdetr-nano-macos.aimodel    (~52 MB, GPU)
uv run rfdetr-convert --variant nano --platform macos --batch 8       # -> rfdetr-nano-macos-b8.aimodel  (GPU throughput)
uv run rfdetr-convert --variant nano --platform ios --profile tiny    # -> rfdetr-nano-ios-tiny.aimodel  (~22 MB)

Other flags: --variant {nano,small,medium,base,large}, --weights <custom.pth>, --num-classes N, --resolution R, --output-dir DIR, --overwrite.

Verify against the PyTorch source (downloads a COCO sample image on first run; loads with the platform's preferred compute unit):

uv run rfdetr-verify --variant nano --platform ios                  # ANE-preferred
uv run rfdetr-verify --variant nano --platform macos --batch 8       # GPU-preferred, batch 8
uv run rfdetr-verify --variant nano --asset path/to.aimodel --image my.jpg

Compile — ahead-of-time (optional)

rfdetr-compile wraps xcrun coreai-build compile, which produces architecture-specific .aimodelc assets that cut on-device first-run specialization time. It maps --platform to coreai-build's --platform/--preferred-compute (iOS→Neural Engine, macOS→GPU):

uv run rfdetr-compile --variant nano --platform ios                  # -> coreai-build ... --platform iOS
uv run rfdetr-compile --variant nano --platform macos --batch 8       # -> ... --platform macOS
uv run rfdetr-compile --variant nano --platform ios --architecture h18p --dry-run

coreai-build ships with a recent Xcode (macOS/iOS 27 SDK) and isn't present in this Xcode, so the command prints the exact invocation it would run and exits cleanly (deferred) — same pattern as the runtime checks. It runs for real once the Core AI command-line tools are installed. Example deferred output:

[COMPILE deferred: coreai-build not found — needs Xcode with the Core AI tools (macOS 27)]
Would run:
  xcrun coreai-build compile exports/rfdetr-nano-ios-small.aimodel --platform iOS \
      --preferred-compute neural-engine --output compiled

Verification results (RF-DETR-Nano, COCO sample 000000039769.jpg)

Every asset detects the same 5 objects as the fp32 PyTorch model — 2 cats, 2 remotes, 1 couch — classes matching. Detection agreement (the primary gate) per platform/profile:

Asset	Compute unit	Detections	Max box diff	Result
rfdetr-nano-ios-small (8-bit palett, b1)	Neural Engine	5 = 5	0.75 px	PASS
rfdetr-nano-macos (fp16, b1)	GPU	5 = 5	0.70 px	PASS
rfdetr-nano-macos-b4 (fp16, b4)	GPU	5 = 5	0.70 px	PASS

Earlier profile sweep (compression only, batch 1): max-accuracy 103 MB / 0.00 px · balanced 52 MB / 0.70 px · small 27 MB / 0.75 px · tiny 22 MB / 1.02 px — all PASS.

On metrics: raw-tensor PSNR is informational only for this model. The dets/labels tensors hold 300 queries but only ~5 are real detections — low-bit rounding on the ~295 unused queries tanks the PSNR while leaving the actual detections essentially unchanged. The primary gate is therefore decoded-detection agreement (count, class, sub-pixel box, score), which all profiles pass.

Compression exploration

The profile ladder comes from a coreai-opt weight-compression sweep (tools/explore_compression.py, using the model-compression-exploration skill's tested helpers). Results: docs/compression_exploration.md · scatter.

Key findings on RF-DETR-Nano: palettization beats quantization at equal size (8-bit palett = 0.75 px drift vs 8-bit quant = 3.77 px); 4-bit breaks detection entirely (wrong classes, boxes off by ~360 px) — so the usable frontier is fp16 → 8-bit → 6-bit palettization. Re-run with uv run python tools/explore_compression.py.

Mixed precision (encoder 4-bit / decoder 8-bit) — tested, does not help here

A common rescue for 4-bit is to compress only the large backbone hard and keep the sensitive decoder at higher precision. The backbone is 76 % of RF-DETR-Nano's params, so this should capture most of the size win. We tested it (palett_mixed_enc4_dec8): backbone → 4-bit, transformer/decoder/heads → 8-bit. It still breaks (wrong classes, ~390 px drift), and finer group sizes (gs8, gs4) don't fix it — the DINOv2 backbone itself can't tolerate 4-bit (only 16 LUT centroids), independent of the decoder. So for this model the encoder is 4-bit-sensitive too, and 6-bit is the practical floor. Pinpointing which individual layers (if any) tolerate 4-bit needs the per-layer activation-PSNR sensitivity analysis (see Roadmap) — a blanket encoder/decoder split is too coarse.

Runtime note: Core AI assets declare minimum_os = v27. This machine runs macOS 26.5, yet the coreai.runtime Python runtime loaded and executed the assets here, so verification ran for real (it was expected to be deferred). Full on-device behaviour (Neural Engine / GPU placement, AOT compilation) should still be confirmed on macOS 27 hardware.

Project layout

src/rfdetr_coreai/
  convert.py        # rfdetr-convert  — RF-DETR → .aimodel (platform + profile + batch)
  compile.py        # rfdetr-compile  — AOT compile via coreai-build (COMPILE step)
  verify.py         # rfdetr-verify   — PyTorch vs Core AI, per-platform compute unit
  profiles.py       # profiles, platform defaults, platform→compute-unit mapping
  decode.py         # detection decode (per-class sigmoid, cxcywh→xyxy)
  export_compat.py  # torch.export shims (assert no-op + manual grid_sample)
tools/explore_compression.py  # reproducible compression sweep
docs/compression_exploration.{md,png,jsonl}  # sweep results
exports/            # generated .aimodel assets   (gitignored)
samples/            # downloaded test images       (gitignored)

Roadmap

• Deeper ANE authoring (needs macOS 27 to measure the payoff): BC1S layout + Conv2d(1×1) rewrite of the DINOv2 ViT (157 Linear layers); the deformable-attention gather stays GPU/CPU.
• Run AOT compile for real: rfdetr-compile is wired but coreai-build is absent in this Xcode; execute it once the Core AI command-line tools (macOS 27 SDK) are installed.
• True dynamic batch: surgery in the deformable decoder's tensor-derived value.split(...) so the batch dim stays symbolic (currently each asset is a fixed batch size).
• Deeper compression: per-layer activation-PSNR sensitivity to find any 4-bit-tolerant layers (the coarse encoder/decoder split didn't work — see Compression exploration).
• Other variants (Small/Medium/Base/Large) and the segmentation/keypoint heads.
• A minimal Swift on-device demo.