Deployment Workflow: PyTorch → Core ML → ANE

Part of the ANE Knowledge Base. Sources: both Apple articles and the export code in references/ml-ane-transformers (README tutorial) and references/ml-vision-transformers-ane/export.py.

The pipeline is identical for text and vision models:

PyTorch model (ANE-optimized, eval mode)
   → torch.jit.trace                      (TorchScript)
   → coremltools ct.convert(mlprogram)    (.mlpackage)
   → Xcode Performance tab                (verify ANE dispatch + measure)
   → ship in app, load asynchronously

1. Prepare the model#

• Build/instantiate the ANE-optimized variant and load the pretrained weights (state-dict hooks handle Linear→Conv2d and LayerNorm order — see doc 03).
• model.eval() always; these ports are inference-only.
• For HF models: load the baseline with return_dict=False and torchscript=True (coremltools does not support dict outputs).
• Wrap tracing in torch.no_grad() (as export.py does) to avoid autograd baggage.

2. Trace to TorchScript#

# Text (fixed sequence length — pad to max_length)
tokenized = tokenizer(["sample"], return_tensors="pt",
                      max_length=128, padding="max_length")
traced = torch.jit.trace(model, (tokenized["input_ids"], tokenized["attention_mask"]))

# Vision (fixed input shape)
x = torch.rand((1, 3, 256, 256))
traced = torch.jit.trace(model, (x,))

Tracing bakes in static shapes — pick the sequence length / resolution / batch size you will serve (Apple's artifacts encode them in the filename: ..._seqLen128_batchSize1, ..._batch1_256x256_...). Export one package per shape you need.

3. Convert with coremltools#

import coremltools as ct, numpy as np

mlpackage = ct.convert(
    traced,
    convert_to="mlprogram",                      # ML Program, not neuralnetwork
    inputs=[
        # text: token ids are int32
        ct.TensorType("input_ids", shape=(1, 128), dtype=np.int32),
        ct.TensorType("attention_mask", shape=(1, 128), dtype=np.int32),
        # vision: ct.TensorType("x", shape=x.shape)
    ],
    compute_units=ct.ComputeUnit.ALL,            # allow CPU+GPU+ANE (default)
)
mlpackage.save("model.mlpackage")

Key points:

• convert_to="mlprogram" targets the modern ML Program backend (FP16 by default on ANE).
• compute_units=ct.ComputeUnit.ALL lets Core ML build the hybrid plan; use CPU_AND_NE during debugging if you want to detect GPU fallbacks.
• Integer inputs (token ids, masks) are declared np.int32.

4. Verify on device: Xcode performance reports#

1. Add the .mlpackage as a resource in any Xcode project.
2. Open it → Performance tab → generate a report on a locally available device (the Mac itself or a connected iPhone/iPad).
3. The report shows per-op compute-unit dispatch (CPU / GPU / ANE) and measured latency.

How to read it:

• Expect the transformer trunk on ANE. Some CPU ops are normal — DistilBERT keeps 4/606 ops (embedding lookups) on CPU because that's genuinely faster.
• Op count grows with chunking → higher one-time load/compile time; steady-state latency is what improves. Load the model asynchronously at app start.
• Re-run reports across target devices/OS versions — dispatch decisions can differ by chip and OS (Apple published curves for iPhone 12/13, iOS 15/16, M1 macOS 13).
• If latency is flat while you shrink the workload, you are bandwidth-bound → consider bigger batches or quantization/pruning (Principle 4).

5. Common pitfalls checklist#

Symptom	Likely cause	Fix
Ops falling back to GPU/CPU mid-graph	Unsupported op or hostile layout (rank > 5, small last axis)	5D relay reshapes; BC1S/NHWC layouts; check each new op in the report
Wildly wrong outputs vs PyTorch	FP16 (eps too small, mask −inf/1e9, overflow)	eps ≥ 1e-7, masks −1e4, optional clamping
Conversion error about dict outputs	HF return_dict=True	return_dict=False (+ torchscript=True)
Checkpoint won't load into optimized model	Linear vs Conv2d weight shapes; LayerNorm scale/bias order	state-dict pre-hooks (linear_to_conv2d_map, bias/weight correction)
Great latency, slow first load	100s of chunked ops compiling	Accept as one-time cost; async load; cache compiled model
Latency doesn't drop with shorter sequences	Bandwidth-bound regime	Increase batch; quantize/prune weights
Padding blow-up / 32× slowdown	Small or singleton last axis	Reorder dims so a large (64-byte-aligned) axis is last; channels multiple of 32

6. Reproducing Apple's exports from the cloned repos#

# Vision: exports tiny-moat-0 at 512x512 and 256x256, global and local attention
cd references/ml-vision-transformers-ane
pip install torch coremltools pytest timm
python export.py          # writes ./exported_model/*.mlpackage
pytest tests.py           # unit tests / usage examples

# Text: follow the README tutorial
cd ../ml-ane-transformers
pip install ane_transformers   # or: pip install -e .
# then run the DistilBERT tutorial code from README.md

That closes the loop: 01 hardware constraints → 02 principles → 03/04 applied architectures → this deployment/verification workflow.