Hardware-aware NAS

"You asked me to minimize FLOPs, so I did. The chip ran it slower. Turns out the chip never read the same textbook I did."

A Chastened AI Agent

Why does a network with half the floating-point operations of another sometimes run slower on your edge accelerator? Because the chip does not bill you by the FLOP. It bills you for memory movement, for kernels its compiler cannot fuse, for depthwise convolutions its neural processing unit (NPU) handles at a fraction of peak throughput, and for activation tensors that spill out of on-chip SRAM. This section assumes you are comfortable with the device taxonomy of Section 59.2 (CPU, GPU, NPU, DSP, MCU) and the memory, power, and thermal budgets of Section 59.3, and that you have met the six-constraint tension web from Chapter 1. The single idea to hold onto: an architecture's cost is a property of the hardware, not of the math, so the search must measure it there. In short: stop optimizing a proxy for speed and optimize speed itself, measured on the chip you will ship.

Why hardware-aware, and what "aware" means

What NAS automates: the choice of architecture itself, kernel sizes, channel widths, block types, depth, expansion ratios, treated as a search over a discrete space rather than a design a human hand-tunes. Why make it hardware-aware: classical NAS optimized accuracy alone and produced networks too heavy for any sensor node. Hardware-aware NAS reframes the search as multi-objective. Formally, for a candidate architecture \(\alpha\) drawn from a search space \(\mathcal{A}\), we want

which the seminal MnasNet work relaxes into a single differentiable-in-spirit reward using a soft penalty,

so that missing the latency target \(\tau\) is punished smoothly rather than by a hard cutoff. How "aware" is achieved: the \(\text{Latency}(\alpha)\) term is not estimated from FLOPs. It is either measured by running the candidate on the physical device, or, far more cheaply, read from a latency lookup table (LUT) that was profiled once per operator per device. When you need it: whenever the deployment target is fixed and unusual enough that its cost curve does not match a desktop GPU, which is nearly every sensor node.

The cost model: latency lookup tables and measured energy

Running every candidate on-device is accurate but glacial: a search may evaluate tens of thousands of architectures. The workhorse compromise is the operator-wise latency LUT. You profile each primitive block (a \(3\times3\) depthwise convolution at a given input resolution and channel count, say) once on the target, store its measured latency, and then estimate any architecture's latency as the sum over its blocks. This assumes block latencies are additive, which holds well on hardware without aggressive cross-layer fusion and is the basis of ProxylessNAS and FBNet. The code below builds a tiny LUT and scores candidate architectures against a latency budget, exactly the inner loop of a hardware-aware search.

import numpy as np
rng = np.random.default_rng(0)

# Latency LUT: measured ms per block type on a target NPU (profiled once).
LUT_MS = {"conv3x3": 0.90, "dwconv3x3": 0.28, "dwconv5x5": 0.41,
          "mbconv6":  0.62, "skip":     0.00}
# Accuracy contribution proxy per block (from a trained supernet).
ACC = {"conv3x3": 0.040, "dwconv3x3": 0.021, "dwconv5x5": 0.028,
       "mbconv6":  0.035, "skip":     0.000}

BLOCKS, DEPTH, BUDGET_MS = list(LUT_MS), 8, 3.0

def sample_arch():
    return [rng.choice(BLOCKS) for _ in range(DEPTH)]

def latency(arch):  return sum(LUT_MS[b] for b in arch)   # additive LUT model
def acc_proxy(arch): return 0.60 + sum(ACC[b] for b in arch)

# MnasNet-style reward: accuracy scaled by a soft latency penalty (w < 0).
def reward(arch, w=-0.07):
    return acc_proxy(arch) * (latency(arch) / BUDGET_MS) ** w

best, best_r = None, -1.0
for _ in range(4000):                       # random search over the space
    a = sample_arch(); r = reward(a)
    if r > best_r and latency(a) <= BUDGET_MS:   # respect the hard budget too
        best, best_r = a, r
print(f"latency={latency(best):.2f} ms  acc_proxy={acc_proxy(best):.3f}")
print("arch:", " -> ".join(best))

Listing 1 is a caricature of the real thing (random search, a hand-set accuracy proxy), but its structure is the real thing. Change one number in LUT_MS, the profile of a different chip, and the search returns a different network. That is the whole point: the architecture is a function of the hardware. Energy and peak-RAM cost models slot into the same objective; for a microcontroller (MCU) the binding constraint is often not latency at all but the largest activation tensor, which must fit in tens of kilobytes of SRAM.

Search strategy: from thousands of trainings to one supernet

Early NAS trained every candidate from scratch and cost thousands of GPU-days, which no sensor-AI team can afford. The breakthrough that made hardware-aware NAS practical is weight sharing. You train a single over-parameterized supernet that contains every candidate architecture as a sub-path; each candidate inherits the supernet's weights instead of training its own. Evaluating an architecture becomes a forward pass, not a training run. Two families dominate. Differentiable NAS (DARTS, FBNet) places a continuous, learnable weight on each candidate operation and adds the LUT latency as a differentiable term in the loss, so gradient descent tunes architecture and weights jointly:

where \(p_{l,o}\) is the (softmaxed) probability of choosing operation \(o\) at layer \(l\) and \(\lambda\) sets the latency pressure. Once-for-All (OFA) decouples training from search entirely: it trains one supernet with progressive shrinking so that every sub-network is independently deployable, then, for each new device, it searches the already-trained supernet against that device's LUT in minutes, no retraining. That "train once, specialize many" property is what lets a single model family target a phone, a smartwatch, and an MCU from one training run.

from ofa.model_zoo import ofa_net
from ofa.nas.accuracy_predictor import AccuracyPredictor
from ofa.nas.search_algorithm import EvolutionFinder

ofa = ofa_net("ofa_mbv3_d234_e346_k357_w1.0", pretrained=True)
finder = EvolutionFinder(AccuracyPredictor(ofa), efficiency_predictor="note10")
best_arch, _ = finder.run_evolution_search(latency_budget=25)   # ms on the target
subnet = ofa.get_active_subnet(best_arch)   # deployable network, no retraining

Pitfalls: the LUT can lie too

What goes wrong: a hardware-aware search is only as honest as its cost model. The additive-LUT assumption breaks when a compiler fuses adjacent operators (so the sum overcounts) or when memory-bandwidth contention makes a block slower in context than in isolation (so the sum undercounts). Why it matters: a mis-calibrated LUT quietly steers the search toward architectures that look fast on paper and stall on-device, reintroducing the very proxy-gap NAS was meant to close. How to guard against it: validate the search's predicted latency against a handful of real on-device measurements of the final candidates before trusting the ranking, and re-profile the LUT whenever the compiler, runtime, or driver version changes. When this bites hardest: on fused NPU toolchains and on shared-memory MCUs, precisely the targets sensor AI cares about most. And as always in this book, the accuracy numbers that feed the search must come from a leakage-safe split, or the search will happily optimize an architecture for a benchmark you cannot reproduce in the field.

What's Next

In Section 59.7, we zoom out from the single device to the full pipeline: how to decide which computation stays on the node and which crosses to the cloud, and how quantization, pruning, and the hardware-aware architectures of this section fit into a repeatable edge-cloud deployment lifecycle rather than a one-off optimization.