Learn AI Series (#122) - Edge AI: Running Models on Devices

avatar

Learn AI Series (#122) - Edge AI: Running Models on Devices

variant-b-12-green.png

What will I learn

  • You will learn why we push AI off the server and onto the device: latency, privacy, reliability, and cold hard cost;
  • what the edge hardware zoo actually looks like, from a Raspberry Pi down to a microcontroller with less memory than a single photo;
  • how depthwise separable convolutions buy you real-time inference on a phone-class chip;
  • the three runtimes that matter -- TensorFlow Lite, ONNX Runtime Mobile, and Core ML -- and which one fits which target;
  • how to run a model in the browser with WebAssembly and WebGPU, zero server required;
  • and a full Raspberry Pi camera-to-prediction pipeline, plus the power and thermal walls that dictate every choice you make out here.

Requirements

  • A working modern computer running macOS, Windows or Ubuntu;
  • An installed Python 3(.11+) distribution with PyTorch and onnxruntime (pip install onnxruntime covers the inference side);
  • You've been through episode #120 (making models fast) and episode #121 (serving them) -- today we take that fast, well-served model and rip the server out from under it.

Difficulty

  • Beginner

Curriculum (of the Learn AI Series):

Learn AI Series (#122) - Edge AI: Running Models on Devices

Last episode I left you with a needle of a question. We had just built a proper serving layer -- batched, cached, autoscaled, safe to roll out -- and then I pointed at the one assumption sitting underneath the whole thing: that your model lives on a fat server with a GPU humming away next to it. But what happens when it doesn't? When the model has to run on the phone in your pocket, offline, on a battery that actively resents you for asking? That is today. Welcome to the edge.

And I want to be very clear about something up front, because it's the mistake almost everybody makes when they first come out here. Edge AI is not "cloud AI, but smaller". It is a genuinely different design paradigm. Different constraints, different tools, different instincts. On the server you optimise for throughput and you have effectively infinite power and memory. On the edge you optimise for fitting at all, and every milliwatt is a decision you have to defend. Billions of devices live in this world already -- phones, doorbells, cameras, drones, cars, insulin pumps, factory sensors -- and none of them get to phone home for every little inference. So let's learn how to think like they do.

Why the edge? Four honest reasons

Whenever someone tells you to move inference onto the device, they're chasing at least one of these four things. Usually more.

Latency. A round trip to the cloud is 50 to 200 milliseconds on a good day, before the model has computed anything at all. An autonomous vehicle deciding whether to brake does not have 200ms to spend on network overhead -- it needs an answer in ten. Keyboard next-word prediction needs to feel instant, which means sub-5ms. Some problems simply cannot tolerate a network hop, full stop.

Privacy. Medical images, fingerprints, faces, the private conversation you're dictating. If the inference happens on the device, the raw data never leaves it. This isn't just a nice-to-have any more -- it's the law in a lot of places (think GDPR, HIPAA). "We never upload your data" is a much easier promise to keep when you literally never upload it.

Reliability. Edge devices have to keep working when the network doesn't. A factory robot can't down tools because the WiFi hiccuped. A medical monitor can't skip an alarm because a server three timezones away is having a bad night. Offline-first isn't a feature here, it's the baseline expectation.

Cost. At scale, edge inference is almost free. A five-dollar microcontroller running a model costs you exactly nothing per prediction after you've shipped it. That same workload on a cloud GPU at a fraction of a cent per call sounds trivial -- until you multiply by a hundred million calls a day and your finance team comes looking for you ;-)

The edge hardware zoo

Before you pick a model you have to know what you're deploying onto, and the range is genuinely wild. "Edge" spans four orders of magnitude of compute.

# A rough field guide to what "the edge" actually means
edge_devices = {
    'raspberry_pi_5': {
        'cpu': 'ARM Cortex-A76 (4 cores, 2.4 GHz)',
        'ram': '4-8 GB',
        'gpu': 'VideoCore VII (not really useful for ML)',
        'npu': None,
        'power': '5-12W',
        'model_budget': '<100MB, INT8 preferred',
    },
    'nvidia_jetson_orin_nano': {
        'cpu': 'ARM Cortex-A78AE (6 cores)',
        'ram': '8 GB unified',
        'gpu': 'Ampere (1024 CUDA cores)',
        'npu': None,
        'power': '7-15W',
        'model_budget': '<500MB, FP16/INT8',
    },
    'modern_smartphone': {
        'cpu': 'ARM Cortex-A/X (8 cores)',
        'ram': '6-12 GB',
        'gpu': 'Adreno / Mali / Apple GPU',
        'npu': 'Dedicated NPU (15-30 TOPS)',
        'power': '3-8W',
        'model_budget': '<200MB for a good user experience',
    },
    'microcontroller': {
        'cpu': 'ARM Cortex-M (1 core, 100-400 MHz)',
        'ram': '256 KB - 1 MB',
        'gpu': None,
        'npu': None,
        'power': '0.01-0.5W',
        'model_budget': '<100KB, INT8 only',
    },
}

Look at that microcontroller row and let it sink in. 256 kilobytes of RAM. That is smaller than most JPEGs on your phone. And people run keyword-spotting models ("Hey Siri", "OK Google") on hardware exactly like that, sipping microwatts, all day. The model budget line is the one that matters, and it is not just "does it fit in RAM". It's about whether inference completes fast enough inside the power envelope. A 100MB model might run in 200ms on a Pi while pulling 10W -- fine for a plugged-in gadget, a disaster for something on a coin cell. Same model on the microcontroller? Not in a hundred years.

The workhorse: depthwise separable convolutions

Here's the single most important architectural trick in edge computer vision, and once you see it you can't unsee it. Back in episode #45 we built ordinary convolutions from scratch, and they are expensive -- every output channel looks at every input channel. On the edge that's a luxury you can't afford, so we split the work in two: a depthwise step that filters each channel on its own, and a pointwise 1x1 convolution that mixes the channels afterward. This is the beating heart of MobileNet and EfficientNet.

import torch
import torch.nn as nn

# A small, deliberately edge-friendly classifier
class EdgeClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),         # a normal first conv
            nn.BatchNorm2d(16),
            nn.ReLU6(),                              # ReLU6 quantizes nicely
            # --- depthwise separable block ---
            nn.Conv2d(16, 16, 3, padding=1, groups=16),  # depthwise: per-channel
            nn.BatchNorm2d(16),
            nn.ReLU6(),
            nn.Conv2d(16, 32, 1),                    # pointwise: mix channels
            nn.BatchNorm2d(32),
            nn.ReLU6(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(32, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = x.flatten(1)
        return self.classifier(x)

model = EdgeClassifier()
model.eval()
print(sum(p.numel() for p in model.parameters()), "parameters")

Why bother? Do the arithmetic, because the arithmetic is the whole argument. A standard 3x3 convolution turning 16 channels into 16 channels costs 16 x 16 x 3 x 3 = 2304 multiplications per spatial position. The depthwise-separable version doing roughly the same job costs 16 x 3 x 3 (depthwise) plus 16 x 32 (pointwise) = 144 + 512 = 656. That's a 3.5x reduction in multiplies, and the gap only widens as the channel counts grow. On a server you'd shrug. On an ARM core trying to hit 30 frames a second, that 3.5x is exactly the difference between "smooth" and "slideshow". Notice I also used ReLU6 instead of plain ReLU -- capping activations at 6 keeps the numeric range tidy, which matters a lot when we quantize to INT8 in a moment (episode #120 territory).

Getting there: PyTorch to TensorFlow Lite

TFLite is still the most widely deployed on-device runtime on the planet, largely because it rides along on Android. But we trained in PyTorch, and PyTorch doesn't speak TFLite directly, so we hop through ONNX as a neutral middle format (we met ONNX back in episode #120).

# Step 1: export the PyTorch model to ONNX
dummy = torch.randn(1, 3, 32, 32)
torch.onnx.export(
    model, dummy, "edge_model.onnx",
    input_names=["input"], output_names=["output"],
    opset_version=13,
)

# Step 2 + 3 happen at the shell, not in Python:
#
#   $ pip install onnx2tf
#   $ onnx2tf -i edge_model.onnx -o tflite_output           # ONNX -> TFLite
#   $ onnx2tf -i edge_model.onnx -o tflite_output -oiqt      # ...with INT8 quantization
#
# The -oiqt flag does full integer quantization during conversion,
# which is where most of your size and speed win actually comes from.

Two hops, one gotcha. The gotcha is that not every exotic PyTorch operation survives the trip cleanly -- if you built something weird, the converter will let you know in the least helpful way possible. Having said that, for the bread-and-butter conv/linear/activation stack that ninety percent of edge models are made of, this path just works. And that -oiqt INT8 flag is not optional flair: quantizing to 8-bit integers typically cuts model size by 4x and speeds inference up by 2 to 4x, because integer math is cheaper and moves less memory around. On the edge, moving less memory is often the whole game.

ONNX Runtime Mobile

You don't actually have to convert to TFLite at all. ONNX Runtime ships a mobile-flavoured build that's lean, ARM-optimised, and -- crucially -- lets you keep the exact same .onnx file across your server and your device. One artifact, many targets. That consistency is worth a lot when you're debugging a discrepancy at 2 AM.

import onnxruntime as ort
import numpy as np
import time

# Tune the session for a small device
options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
options.intra_op_num_threads = 2   # don't fight the OS for all the cores

session = ort.InferenceSession("edge_model.onnx", options)

x = np.random.randn(1, 3, 32, 32).astype(np.float32)
result = session.run(None, {"input": x})
print("Predicted class:", int(np.argmax(result[0])))

# Always, always benchmark on the actual device
times = []
for _ in range(100):
    start = time.perf_counter()
    session.run(None, {"input": x})
    times.append(time.perf_counter() - start)
print(f"Median latency: {np.median(times) * 1000:.1f} ms")

The same Python API you see here exists in C++, Java, Swift and C# -- so you prototype on your laptop and ship the identical logic into an Android or iOS app with barely a change in shape. That intra_op_num_threads knob is a small thing that bites people: crank it to all cores and you'll actually get worse latency on a phone, because you're now brawling with the operating system and every other app for the same silicon. Leave it a bit of breathing room.

Core ML: talking to Apple's Neural Engine

If your target is an iPhone, iPad or a Mac, the runtime you want is Core ML, because it's the only one that can hand work to Apple's dedicated Neural Engine (the ANE). That thing runs INT8 and FP16 inference absurdly efficiently -- frequently faster than the GPU and at a sliver of the power.

import coremltools as ct

traced = torch.jit.trace(model, torch.randn(1, 3, 32, 32))

mlmodel = ct.convert(
    traced,
    inputs=[ct.TensorType(name="input", shape=(1, 3, 32, 32))],
    compute_precision=ct.precision.FLOAT16,   # let the Neural Engine shine
)
mlmodel.save("EdgeClassifier.mlpackage")

# On the device this becomes a couple of lines of Swift:
#
#   let model = try EdgeClassifier(configuration: MLModelConfiguration())
#   let output = try model.prediction(input: preprocessed)

The bit I find genuinely clever is that Core ML does the hardware routing for you, per layer. It'll quietly run one operation on the ANE, the next on the GPU, and a stubborn third on the CPU, all inside a single forward pass, picking whatever is cheapest for each. You don't orchestrate any of it. From your side it's one prediction() call, and Apple's compiler has already fought the placement battle on your behalf. Nota bene: you can only build Core ML models on a Mac, so this path is a hard no if your dev machine runs Windows or Linux.

WebAssembly: AI right there in the browser

Here's an edge target people forget is even an edge target: the browser tab. Run the model in WebAssembly (with WebGPU for acceleration where it exists) and you get zero server cost, zero upload of user data, and inference that starts the instant the page loads.

# This is JavaScript running in the browser -- shown as a string
# so we can talk about it without leaving Python for the walkthrough.
web_inference = """
import * as ort from 'onnxruntime-web';

async function runModel(inputData) {
    // Model is cached by the browser after the first fetch
    const session = await ort.InferenceSession.create('model.onnx', {
        executionProviders: ['webgpu', 'wasm'],   // WebGPU first, WASM fallback
    });

    const tensor = new ort.Tensor('float32', inputData, [1, 784]);
    const results = await session.run({ input: tensor });
    return Array.from(results.output.data);
}
"""
print(web_inference)

WebGPU is the new browser standard for talking to the GPU, and with ONNX Runtime Web sitting on top of it you can get near-native inference speed straight inside Chrome or Firefox. For older browsers without WebGPU it quietly falls back to plain WASM -- slower, but it runs literally everywhere. What can you build with this? Real-time image filters, on-device OCR, text classification that never phones home, even speech recognition with a Whisper-class model running fully local (a distant cousin of the local-models discussion from episode #70). No install, no account, no data leaving the machine. That last part is a privacy story you can actually put on a landing page, and mean it ;-)

A real Raspberry Pi pipeline

Enough fragments -- let's wire an end-to-end thing you could genuinely run this weekend. Camera in, prediction out, on a Pi.

# deploy_to_pi.py -- runs on the Raspberry Pi itself
import onnxruntime as ort
import numpy as np
import time
from picamera2 import Picamera2

class EdgeInferencePipeline:
    def __init__(self, model_path, labels_path):
        options = ort.SessionOptions()
        options.intra_op_num_threads = 4    # the Pi 5 has 4 cores, use them here
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        self.session = ort.InferenceSession(model_path, options)
        self.input_name = self.session.get_inputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape
        with open(labels_path) as f:
            self.labels = [line.strip() for line in f]

    def preprocess(self, image):
        """Resize + normalise into the model's expected NCHW tensor."""
        h, w = self.input_shape[2], self.input_shape[3]
        image = np.array(image)[:h, :w, :3]
        image = image.astype(np.float32) / 255.0
        image = np.transpose(image, (2, 0, 1))       # HWC -> CHW
        return np.expand_dims(image, 0)

    def predict(self, image):
        x = self.preprocess(image)
        out = self.session.run(None, {self.input_name: x})
        probs = self._softmax(out[0][0])
        top = int(np.argmax(probs))
        return self.labels[top], float(probs[top])

    def _softmax(self, x):
        e = np.exp(x - np.max(x))
        return e / e.sum()

    def run_camera_loop(self, fps_target=5):
        camera = Picamera2()
        camera.start()
        frame_interval = 1.0 / fps_target

        while True:
            start = time.perf_counter()
            frame = camera.capture_array()
            label, confidence = self.predict(frame)
            elapsed = time.perf_counter() - start
            print(f"{label}: {confidence:.1%} ({elapsed * 1000:.0f}ms)")

            # Duty-cycle to hold the target frame rate (and the thermals)
            sleep_time = frame_interval - elapsed
            if sleep_time > 0:
                time.sleep(sleep_time)

On a Raspberry Pi 5 running a quantized MobileNetV2, this classifies frames at roughly 15 to 30 FPS depending on input resolution -- comfortably real-time for object classification, gesture recognition, keyword spotting, and a whole shelf of practical little products. Notice intra_op_num_threads=4 here, the opposite of my advice on the phone: on a dedicated Pi doing nothing but this, you want all four cores. Context decides. And see that little time.sleep at the bottom? That's not laziness, that's duty cycling -- deliberately resting the chip between frames so it doesn't cook itself. Which brings us to the part the tutorials love to skip.

The part nobody warns you about: power and heat

On a server you never think about watts or temperature -- someone else's problem, three datacenters away. On the edge these are your problem, and they will quietly wreck your beautiful model if you ignore them.

edge_reality = """
POWER BUDGET
- Battery devices : 1-5W sustained, brief bursts to ~10W
- Plugged devices : 5-25W
- Rough model draw = operations_per_second x energy_per_operation
- INT8 ops use about 4x LESS energy than FP32 ops

THERMAL
- Passive cooling (no fan) : ~2-5W sustained before it throttles
- Active cooling (tiny fan): ~10-25W sustained
- Sustained heavy inference -> throttling -> slower inference (a spiral)
- Fix: duty cycling -- run 50ms, sleep 50ms, keep the average down

MEMORY
- Pi            : 4-8GB, shared with the whole OS
- Smartphone    : maybe 1-3GB free for your app
- Microcontroller: 256KB, total, for everything
- Peak memory = model weights + LARGEST intermediate activation
- You must budget for the activations, not just the weights
"""
print(edge_reality)

That thermal spiral is the sneaky one. Push a passively-cooled board too hard and it heats up, so the firmware throttles the clock to protect itself, so your inference gets slower, so it takes longer, so it stays hot -- and now your promised 30 FPS is limping along at 12 and you have no idea why. Duty cycling breaks the loop by capping the average load, which is why that unglamorous time.sleep earned its place in the pipeline above. The memory line deserves a second read too: your peak footprint is the weights plus the biggest intermediate activation tensor during the forward pass, not just the file size on disk. Plenty of models that "fit" on paper fall over mid-inference because one fat activation map briefly needs more room than the device has.

The honest summary is this: edge AI is a constraints game, not a research game. The techniques are well understood and mostly settled -- MobileNet, EfficientNet, quantization, pruning, the runtimes we walked through today. The actual work, the part that eats your week, is the unglamorous engineering of making a known-good model run reliably on one specific piece of hardware inside one specific power and heat budget. That's not a knock. That's just where the difficulty genuinely lives out here, and knowing that up front is half the battle.

Let's recap

  • edge AI is its own paradigm, not shrunk-down cloud AI -- you move inference onto the device for lower latency, real privacy, offline reliability, and near-zero marginal cost;
  • the hardware zoo spans four orders of magnitude, from a Jetson with a real GPU down to a microcontroller with 256KB of RAM, and the "model budget" is set by inference-time-within-a-power-envelope, not just fitting in memory;
  • depthwise separable convolutions (MobileNet, EfficientNet) cut the multiplies 3-5x versus ordinary convolutions -- that ratio is what makes real-time inference on an ARM core possible at all;
  • TFLite, ONNX Runtime Mobile, and Core ML are the three runtimes that matter; ONNX Runtime keeps one artifact across server and device, Core ML is your only door to Apple's Neural Engine;
  • WebAssembly + WebGPU turns the browser tab into an edge device -- zero server cost, zero data upload, inference that runs everywhere;
  • and it's a constraints game: INT8 saves ~4x energy, duty cycling defeats thermal throttling, and peak memory is weights plus the largest activation -- budget for all three or the device will remind you the hard way.

Exercises

Three to wrestle with before next time. As always, get your hands dirty first -- I'll walk through full solutions in the following episode.

  1. Export and benchmark. Take the EdgeClassifier above, export it to ONNX, and load it with onnxruntime. Run 100 inferences and report the median latency with intra_op_num_threads=1 versus =4 on your machine. Write one sentence explaining why "more threads" didn't necessarily mean "faster".

  2. Prove the conv savings. Write a small function conv_cost(in_ch, out_ch, k) that returns the multiplies-per-position for a standard convolution, and a second separable_cost(in_ch, out_ch, k) for the depthwise-separable version. Run both for a few channel sizes (16->32, 64->128, 128->256) and confirm the reduction lands in the 3-5x range as the channels grow.

  3. Duty-cycle for a budget. Write a loop that "runs inference" (fake it with a time.sleep of, say, 40ms of work), then sleeps to hold an average load target you pick (e.g. 50% duty cycle). Measure your effective FPS and your actual duty cycle over 10 seconds, and print both. Bonus: what happens to your FPS if the "work" suddenly jumps to 90ms -- does your budget still hold?

We now have a model that can run on a server, get served to the world, and run untethered on a device in someone's hand. But a model that's out there answering real requests raises a nervous question we've been politely avoiding: how do you know it's still working tomorrow, when the data quietly shifts under its feet and nobody's watching? That's where we head next.

Bedankt en tot de volgende keer!

@scipio

stem stemsocial steemstem python programming
0
0
0.000
0 comments