PyTorch Mobile: Lightweight Inference on Embedded Linux

On Cortex-A Linux targets (i.MX8, Jetson Nano/Orin, Raspberry Pi CM4, Snapdragon embedded SKUs), PyTorch models rarely run in their training form. TorchScript export plus the LibTorch mobile/CPU runtime lets you keep the PyTorch operator semantics and autograd-trained weights while dropping the Python interpreter, JIT tracing overhead, and most of the memory footprint that make desktop PyTorch unsuitable for a 512 MB–2 GB embedded target. Understanding what "PyTorch Mobile" actually replaces — and where it still loses to TFLite Micro, ONNX Runtime, or vendor NPU SDKs — is the difference between a inference pipeline that hits its latency budget and one that thrashes swap.

What PyTorch Mobile Actually Is

"PyTorch Mobile" is not a separate framework — it's a deployment path:

1. Model authoring: standard PyTorch (torch.nn.Module), trained on desktop/server with full autograd.
2. Export to TorchScript: either torch.jit.trace() (records ops for a fixed input shape/path) or torch.jit.script() (parses Python control flow into a static IR). Output is a .pt file containing a serialized graph plus weights.
3. Runtime: LibTorch compiled for the target (ARM64, ARMv7, sometimes with XNNPACK or QNNPACK backends), or on Android/iOS the dedicated org.pytorch:pytorch_android_lite / LibTorch-Lite builds. On embedded Linux the relevant build is LibTorch for ARM Linux, invoked via the C++ API (torch::jit::load(...)) or Python if space allows.
4. As of PyTorch 2.x, the newer ExecuTorch runtime is the intended successor for constrained edge/embedded targets, but LibTorch mobile remains the production path for many current Cortex-A Linux deployments.

Key distinction from TFLite Micro: PyTorch Mobile targets Cortex-A + Linux (MMU, filesystem, dynamic memory), not bare-metal Cortex-M. It's the right tool when you already have a Linux BSP and need PyTorch-trained models without a conversion hop through ONNX/TFLite, not for sub-1MB RAM microcontrollers.

Trace vs. Script: Which Export Path

Practical rule: trace CNN backbones (MobileNet, ResNet, most vision models), script anything with sequence loops or shape-dependent logic. Always validate traced output against eager-mode output on several inputs before trusting the export — a mistraced conditional silently ships wrong behavior with no runtime error.

import torch, torchvision

model = torchvision.models.mobilenet_v3_small(weights="DEFAULT").eval()
example = torch.rand(1, 3, 224, 224)

traced = torch.jit.trace(model, example)
traced.save("mobilenet_v3_small.pt")

# sanity check: outputs must match within fp tolerance
with torch.no_grad():
    ref = model(example)
    got = traced(example)
assert torch.allclose(ref, got, atol=1e-5)

Quantization for the CPU Backend

Cortex-A cores without NPU/DSP offload run PyTorch Mobile inference on NEON via the QNNPACK or XNNPACK backend. INT8 quantization is the primary lever for both latency and memory:

• Dynamic quantization (torch.quantization.quantize_dynamic): weights INT8, activations quantized on the fly at inference. Simple, no calibration data needed, best for linear/LSTM-heavy models. Little benefit for conv-heavy vision models since activations dominate compute.
• Static (post-training) quantization: weights and activations both INT8, activation ranges fixed via a calibration pass over representative data. Requires inserting QuantStub/DeQuantStub, fusing conv-bn-relu, then running calibration. Gives the largest speedup for CNNs — QNNPACK's INT8 conv kernels are roughly 2–4× faster than FP32 on ARM NEON, plus 4× memory reduction for weights.
• Quantization-aware training (QAT): same graph transform as static PTQ, but fake-quant ops are active during fine-tuning, recovering accuracy lost to naive PTQ — typically 1–3 percentage points on classification tasks with aggressive INT8.

model.qconfig = torch.quantization.get_default_qconfig("qnnpack")
torch.quantization.prepare(model, inplace=True)
# run calibration_loader through model here (no grad, no backward)
torch.quantization.convert(model, inplace=True)
torch.backends.quantized.engine = "qnnpack"

Worked Example: Latency Budget on a Cortex-A53 Quad-Core

Target: MobileNetV3-Small (≈2.9M params, ≈56M MACs at 224×224) on a 1.4 GHz Cortex-A53 quad-core (e.g., typical i.MX8M-class SoC), single-threaded inference, no NPU.

FP32 baseline estimate:

• A53 sustains roughly 2–4 FP32 GFLOPS/core in practice for conv-heavy workloads (well below theoretical NEON peak due to memory bandwidth and instruction overhead).
• Model needs ≈2 × 56M = 112M FLOPs (MAC = 2 FLOPs).
• Estimated time ≈ 112M / 3G ≈ 37 ms per inference, single core.

INT8 static-quantized estimate:

• QNNPACK INT8 kernels typically deliver 2–3× throughput over FP32 on the same core for 3×3/1×1 convs dominant in MobileNet.
• Estimated time ≈ 37 ms / 2.5 ≈ 15 ms.

Check against a real constraint: if the target is 30 fps video classification (33 ms/frame budget), FP32 single-core (37 ms) already misses it; INT8 (15 ms) leaves headroom for pre/post-processing (resize, normalize, softmax, NMS if detection) and other frame-pipeline work. This is why quantization isn't optional polish here — it's the difference between meeting and missing the real-time budget on a mid-range Cortex-A53.

Memory check: FP32 weights ≈ 2.9M × 4 B ≈ 11.6 MB; INT8 weights ≈ 2.9M × 1 B ≈ 2.9 MB, plus small per-channel scale/zero-point tables. On a 512 MB Linux target running other services, this reduction matters for resident set size, not just cache behavior.

Threading and Memory Practicalities

• Thread count: LibTorch defaults to using all available cores via its intra-op thread pool (at::set_num_threads). On a shared SoC also running camera capture, display compositing, and other services, pin thread count explicitly (often 2 of 4 cores) rather than letting inference starve the rest of the system.
• Memory allocator: LibTorch's default CPU allocator can fragment under repeated alloc/free of intermediate tensors. For steady-state inference loops, reuse pre-allocated output tensors (torch::jit::IValue reuse, or torch::NoGradGuard plus explicit output buffers) instead of allocating a fresh tensor per call.
• Binary size: a full LibTorch build is 40–100+ MB. For embedded deployment, build with -DBUILD_LITE_INTERPRETER=ON and strip unused operators via the selective build / custom op registration path — this is essential when flash/storage is constrained, and can cut the runtime to under 10 MB depending on operator coverage needed.
• Startup latency: TorchScript deserialization and operator registration add measurable cold-start time (tens to low hundreds of ms) — relevant for systems that spin up inference on demand rather than keeping a resident process.

Where PyTorch Mobile Loses to the Alternatives

PyTorch Mobile's strongest case is when the training and deployment teams are the same, the target is Cortex-A Linux, and re-authoring the model in another framework's ops isn't worth the engineering cost — common in robotics, drones, and industrial vision boxes already running a full Linux BSP.

Key Takeaways

• PyTorch Mobile = TorchScript export (trace or script) + LibTorch CPU runtime (QNNPACK/XNNPACK backend), targeting Cortex-A Linux, not bare-metal MCUs.
• Trace fixed-topology CNNs; script models with data-dependent control flow; always numerically verify the export against eager mode.
• Static INT8 quantization gives the largest CPU speedup (~2–4×) and a 4× weight-memory reduction versus FP32; dynamic quantization mainly helps linear/RNN layers.
• Real latency budgets must be checked against actual core throughput, not FLOPs alone — memory bandwidth and kernel efficiency dominate on Cortex-A53-class cores.
• For NPU-equipped targets or bare-metal MCUs, PyTorch Mobile is the wrong tool — use the vendor SDK/Ethos-U toolchain or TFLite Micro/CMSIS-NN respectively.
• Selective/lite builds and pinned thread counts are necessary for production embedded deployment; default LibTorch builds and thread pools are sized for desktop, not shared-SoC Linux systems.