On-Device and Edge ML Questions

Start with a reproducible test harness and narrow the problem with controlled experiments.

1. Baseline setup

• Run both models (TensorFlow and Core ML) on the exact same inputs and seed any nondeterminism. Use a fixed set of representative tensors (random, edge values, and real samples).
• Log inputs, intermediate activations (where possible), final outputs, and metadata (model versions, conversion flags).

2. Check preprocessing

• Verify the same input scaling, normalization, color ordering (RGB vs BGR), and mean/std subtraction. Compare byte-level input fed to each model.
• Mitigation: Align preprocessing pipeline in code or bake normalization into the model (add a preprocessing layer) so both see identical tensors.

3. Data layout (NHWC vs NCHW)

• Confirm framework and Core ML layer expected layouts. Provide a simple identity test: feed an input with a unique pattern per channel and inspect outputs to detect channel/axis swaps.
• Mitigation: Insert explicit transpose ops during conversion or adjust input layout in client code.

4. Precision changes (FP32 → FP16/quantization)

• Check if conversion enabled float16 or quantization. Run both FP32 and lowered-precision versions and compare deltas.
• Mitigation: Disable precision lowering to confirm cause; if needed, calibrate quantization, increase bit-width, or fine-tune/quantization-aware-train to recover accuracy.

5. Op implementation differences

• Isolate layers with largest error by comparing layerwise activations (export intermediate nodes or replicate subgraphs). Pay attention to ops with approximate implementations (softmax, layer-norm, pooling, custom ops).
• Mitigation: Replace problematic ops with equivalent sequences supported identically in Core ML, implement custom layers (MLCustomLayer) replicating TF numeric behavior, or adjust tolerances.

6. Numerical stability and non-determinism

• Test with deterministic settings, compare results with varying batch sizes and single-sample runs.
• Mitigation: Use stable numerics (add epsilons), ensure same reduction axes and epsilon placement in batchnorm/normalization.

7. Quantify and accept

• Define acceptable numeric thresholds (absolute and relative). For each test, report max, mean absolute/relative differences and mismatch rate on downstream metrics (accuracy, AUC).
• If differences are small and downstream metrics unchanged, document and accept; otherwise iterate on the targeted fixes above.

This systematic isolating approach (input → layout → precision → ops → numeric) lets you pinpoint root cause and apply minimal mitigations.

Requirements:

• Functional: predict on-device with low latency, privacy, offline support.
• Non-functional: small binary size, low memory/CPU, energy-efficient, ability to update/rollback.

High-level lifecycle:

1. Training & export

• Train in PyTorch/TensorFlow; instrument training to produce deterministic checkpoints, export to interoperable formats (TorchScript/ONNX/TF SavedModel).
• Validate outputs on representative mobile inputs and measure latency/accuracy on CPU.

2. Conversion to Core ML (.mlmodel)

• Use coremltools to convert from ONNX/TorchScript/TensorFlow to .mlmodel. Specify input/output shapes, precision (fp32 → fp16/INT8), and metadata (version, signature).
• Run unit tests comparing outputs to reference to ensure numerical parity.

3. Compilation to .mlmodelc

• Compile with coremlcompiler for target iOS version/arch to produce optimized .mlmodelc. This optimizes operator kernels and bundles metadata for runtime.
• Test on physical devices across CPU/GPU/Neural Engine configurations.

4. Packaging & delivery

• Embed in app bundle for simplest, offline-first delivery; good for small stable models and faster startup.
• Use on-demand download (App Store-hosted assets or secure CDN) for large models or frequent updates. Implement integrity checks (signatures, checksums) and atomic swaps.

5. Versioning & rollback

• Embed semantic version + model hash in metadata. Serve version manifest from backend.
• Client-side: keep current and previous model copies; download new model to temp location, validate, then switch atomically. If validation fails or performance regresses, rollback to previous version.
• Use staged rollouts and A/B tests; monitor key metrics and device telemetry.

Key mobile constraints & strategies:

• Size: quantize to fp16/INT8, prune, distill; compress with packaging and lazy-loading.
• Latency/CPU: use Core ML acceleration, operator fusion, and batching where appropriate.
• Memory: stream inputs, limit peak activation memory, prefer smaller architectures.
• Energy: prefer Neural Engine, reduce frequency of inference, cache results.
• Security: sign models, encrypt downloads, respect user privacy.
• Observability: collect lightweight telemetry and simulated benchmarks to detect regressions.

Trade-offs: aggressive quantization reduces size but can reduce accuracy; runtime downloads allow agility but add network/SEC complexity.

Brief approach: allow shape ranges on the Core ML input type so the converted model accepts multiple image resolutions, and tune precision/caching to balance ANE latency and memory.

Example conversion (coremltools):

import coremltools as ct
from coremltools.models.datatypes import TensorType, RangeDim

# PyTorch/TF model -> traced_model
# allow batch=1, channels=3, height/width in [128,1024]
input_type = TensorType(shape=(1, 3, RangeDim(128, 1024), RangeDim(128, 1024)))
mlmodel = ct.convert(traced_model, inputs=[input_type],
                     convert_to="mlprogram",
                     compute_precision=ct.precision.FLOAT16)  # optional
mlmodel.save("flexible_image.mlpackage")

Core concepts / options:

• RangeDim: specify min/max for dynamic H/W.
• TensorType vs ImageType: ImageType provides preprocessing metadata (scale/mean) but TensorType gives explicit dims for flexibility.
• convert_to="mlprogram": ML Program format better supports dynamic shapes than legacy NN.
• compute_precision=FLOAT16: reduces memory and model parameter size; may slightly reduce accuracy.
• add_preprocessing in model vs do it at app: doing preprocessing in-app can avoid recompile when only scale/crop changes.

Runtime trade-offs (ANE and memory):

• ANE supports ML Program and dynamic shapes but shape variability can prevent aggressive static optimizations; the runtime may recompile or reconfigure kernels per new shape, adding one-time or per-shape overhead.
• Larger/resolution inputs increase ANE compute time and peak memory; ANE has limited fast memory — very large images can cause fallback to GPU/CPU or higher latency.
• FLOAT16 reduces memory pressure and can improve ANE throughput, but watch numerical effects.
• Model size on disk: flexible-shape metadata is minimal — parameter count unchanged. Using FLOAT16 reduces disk size.
• Recommendation: constrain ranges tightly (e.g., 224–512) to reduce recompiles, perform deterministic resizing/scaling in the app, and prefer ML Program + float16 for ANE-efficient deployments.

Requirements & goals:

• Three model variants (small/med/large) optimized for different Apple device classes (e.g., A-series low-end, M-series mid, Mx high-end).
• Minimal engineering overhead: reuse pipelines, automated validation, safe rollout, seamless user migration.

High-level strategy:

1. Model build & CI/CD

• Single repo with model recipes + optimization configs (quantization, pruning, compiler flags).
• Parameterized build pipeline (GitHub Actions / Jenkins) that runs a matrix job producing variant artifacts via toolchain (TorchScript/ONNX → Core ML / mlmodelc). Reuse common training artifacts; only run variant-specific optimization steps.
• Artifacts stored in model registry (versioned: semantic + variant tag). CI enforces linters, unit tests, model-size/latency contracts, and signs artifacts.

2. Automated testing on device farms

• Use cloud device farms (AWS Device Farm / private Lab with Xcode bots) and on-prem device pool. CI triggers distributed tests per artifact:
  • Unit inference tests (numerical parity on simulator)
  • Performance tests (latency, memory, peak CPU/GPU, battery impact) on representative physical devices per class
  • Robustness tests (randomized inputs, quantization noise) and golden-output checks for acceptable accuracy deltas
• Test matrix only runs full battery on changed models or when core infra changes; cheap smoke tests otherwise.

3. Runtime model selection logic

• On app startup evaluate device capability profile: CPU/GPU/Neural Engine detection, OS version, available RAM, thermal state, current battery saver mode, and optional user preference.
• Selection policy:
  • Deterministic mapping: device-class → default variant
  • Runtime heuristics: if low battery/thermal throttling, downgrade to smaller variant
  • Graceful fallback: attempt load preferred variant, on OOM/slow load fallback to next smaller
• Implement A/B flags and remote config (feature flagging) to route users to variants for experiments or phased rollouts.

4. Backward compatibility & user migration

• Model contracts: define stable input/output schemas and post-processing compatibility layers so new variants maintain prediction semantics within defined tolerances.
• Store model metadata: version, variant, accuracy/thresholds. On model update, run regression suites; if breaking change, ship compatibility adapter in app that maps between old and new output formats.
• Migration strategies: staged rollout via feature flags (percent rollout), canary on internal users, monitor key metrics (accuracy, latency, crash rates, retention). Provide deterministic seed for experiments to ensure reproducibility.
• User data considerations: if model produces cached user-specific state, include migration code to translate caches or invalidate safely.

5. Observability & rollback

• Telemetry: per-inference latency, memory usage, error counts, model variant tag, and quality signals (if available).
• Automate alerts on SLA breaches; CI/CD ties deployment to metric gates. Fast rollback via remote config toggles.

Trade-offs and engineering efficiency

• Reuse training pipeline and parameterize optimization to avoid duplicate workflows.
• Heavy physical-device tests only on artifacts or infra changes to reduce cost.
• Remote-config driven selection reduces app releases for variant tuning.

This approach balances reliability, low overhead, and safe progressive rollout across the Apple device ecosystem.

Approach summary:

1. Prepare your PyTorch conv model for QAT (module fusion, insert fake-quant modules), 2) run QAT training with observers and learning-rate schedule, 3) convert a trained model to an exportable format (TorchScript preferred), 4) convert to Core ML with coremltools targeting ANE (use float16 where helpful) and watch conversion pitfalls.

Step-by-step with examples

1. Model prep for QAT

• Fuse conv/bn/relu where possible (improves quantization accuracy).
• Replace nn.Modules with quant-aware versions and add FakeQuant/Observers via torch.quantization utilities.

import torch, torch.nn as nn
from torch.quantization import fuse_modules, prepare_qat, convert

class SmallConv(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3,16,3,padding=1)
        self.bn = nn.BatchNorm2d(16)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(16*32*32, 10)
    def forward(self,x):
        x = self.relu(self.bn(self.conv(x)))
        x = x.view(x.size(0),-1)
        return self.fc(x)

model = SmallConv()
model.train()
fuse_modules(model, [['conv','bn','relu']], inplace=True)
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')  # cpu/qconfig choice
prepare_qat(model, inplace=True)

2. QAT training considerations

• Train with lower LR initially (fine-tune from pretrained if available).
• Train long enough for observers to converge (often several epochs).
• Use calibration-like batches (diverse data) before switching to eval for conversion.
• Keep model in train() during QAT steps; switch to eval() and call convert() after training:
  • convert(model.eval(), inplace=False) will produce a quantized model (if using static quant).
• Monitor accuracy vs float baseline; use EMA or regularization if needed.

3. Export path: TorchScript (recommended) vs ONNX

• TorchScript (trace or script) is more stable for PyTorch-specific ops and quantized models.
• For quantized models, tracing a post-conversion model in eval() is recommended:

quantized = convert(model.eval())
scripted = torch.jit.trace(quantized, torch.randn(1,3,32,32))
scripted.save("model_quant.pt")

• ONNX can work but often fails on quantized ops or custom fused ops; if using ONNX, run export on a non-quantized or dequantized model and ensure ops are supported.

4. coremltools conversion notes

• Use coremltools >= 5.x for best PyTorch/TorchScript support.
• Convert from TorchScript when possible:

import coremltools as ct
ts_model = torch.jit.load("model_quant.pt").eval()
example_input = torch.randn(1,3,32,32)
mlmodel = ct.convert(ts_model, inputs=[ct.TensorType(shape=example_input.shape)])
# target ANE: set compute units at runtime when saving/using
mlmodel.save("model.mlpackage")

• To ensure ANE usage at runtime (iOS) specify .mlmodel/ mlpackage compute units in the app: .model.predictionOptions.computeUnits = .all (ANE) or use coremltools.convert(..., compute_units=ct.ComputeUnit.ALL) where supported.
• Consider converting weights to float16 to reduce size and speed on ANE: ct.convert(..., minimum_deployment_target=ct.target.iOS17, compute_precision=ct.precision.FLOAT16)

5. Pitfalls & tips

• Core ML/ANE does not consume PyTorch fake-quant ops—Core ML expects float or float16 networks. True INT8 quantized PyTorch models rarely map directly; prefer converting a dequantized (float) TorchScript model or use weight-only float16 compression.
• Unsupported ops: custom layers, in-place ops, or complex control flow may fail. Replace unsupported ops or provide pure-PyTorch scripted implementations.
• Dynamic shapes: ANE/Core ML prefers static shapes; fix input sizes or provide range constraints.
• Remove training-only modules/observers: ensure model.eval() and convert() before export.
• Validate outputs numerically after each step with representative inputs; compare float baseline vs Core ML outputs within acceptable tolerance.
• If using ONNX: set opset high enough, run onnx-simplifier, and check with onnxruntime before coremltools conversion.

Why this path:

• QAT in PyTorch yields accuracy close to float while simulating quantization during training.
• TorchScript is most compatible with coremltools; Core ML/ANE currently works best with float/float16 models — weight quantization (int8) is not directly supported for execution on ANE in most conversion paths, so use float16 for ANE acceleration when possible.

This process ensures best accuracy retention from QAT training and a reliable conversion path for Core ML targeting the ANE.