Picture a warehouse robot rolling down an aisle at walking pace. A camera on its arm looks at every pallet it passes, and somewhere in that stream of frames there is a box with a crushed corner. The robot has to notice, decide whether it is damaged enough to flag, and do it in the time it takes to drive past, without a network connection to a data centre and on a battery it shares with its motors. That is the situation this article is about.
It has two halves that usually get taught separately, and that is a mistake. The first half is the model: how a Vision Transformer (ViT) turns a picture into a prediction, built from scratch so nothing is magic. The second half is the machine: what happens when you take that model out of a notebook and put it on a Jetson-class board, a phone, or a camera with a small chip in it. The two halves belong together because the architecture decides what is cheap and what is expensive on the device. If you know why attention costs what it costs, you know which knobs to turn when the latency budget is blown.
This is Part 4 of the MLHub series. Part 1 covered agents and MCP, Part 2 covered LoRA and QLoRA, and Part 3 covered Polars. Part 5 will be the mathematical foundations of AI, so here I will use just enough maths to make the mechanisms clear and no more. You should be comfortable reading Python and PyTorch. You do not need to have trained a vision model before.
Here is the plan:
- Why convolutional networks ruled vision for a decade, and what a ViT deliberately throws away.
- The ViT pipeline step by step, with the arithmetic done by hand.
- A NumPy attention function, then a complete ViT in about 100 lines of PyTorch, with shapes printed.
- A parameter and FLOP estimate for ViT-Base.
- The quadratic cost of attention and how the field responded.
- What changes when the model lives in a robot, and an optimisation ladder: smaller model, ONNX, ONNX Runtime, INT8, FP16, TensorRT.
- A latency benchmark that does not lie.
- Production reality: drift, calibration, fallbacks, updates, privacy, and a debugging checklist.
Where I show numbers, they are either computed by hand in the text or clearly labelled as illustrative. I have not measured your hardware, and neither has anyone else. Where I give you a script, run it on your own device.
Why CNNs were king
For roughly ten years, if you wanted to classify images you used a convolutional neural network. AlexNet, VGG, ResNet, EfficientNet: diffrent depths and tricks, same core idea. That idea is worth understanding properly, because a ViT is best explained as "a CNN with the assumptions taken out".
A convolution slides a small filter, say 3x3 pixels, across the image and computes the same weighted sum at every position. Two assumptions are baked into that one operation.
The first is locality. A 3x3 filter only looks at a pixel and its immediate neighbours. The bet is that the useful information for detecting an edge, a corner or a texture lives close together. Stack many layers and the receptive field grows: a unit deep in the network sees a large region, but it built that view up through many small local steps.
The second is translation equivariance. Because the same filter is applied everywhere, if a cat moves ten pixels to the right in the input, its features move ten pixels to the right in the output. The network does not need to learn "cat at the top left" and "cat at the bottom right" as separate facts. A pooling layer at the end then turns equivariance into something closer to invariance: is there a cat anywhere?
Those two assumptions are called inductive biases: things the architecture believes about the world before it has seen a single training example. They are enormously useful. They shrink the number of parameters (one small filter covers the whole image), they make learning fast on modest datasets, and they match how natural images actually behave. That is why CNNs worked so well with a million labelled images, and why they still run happily on phones.
They also have a cost. Locality means that relating two distant parts of an image takes many layers. Weight sharing means the model treats every location the same way, even when position matters. And the biases are fixed: if the data would prefer a different structure, the network cannot choose one.
What ViT drops
In 2020, Dosovitskiy and colleagues published "An Image is Worth 16x16 Words" (arXiv:2010.11929). The idea was almost provocatively simple: take the Transformer encoder from language modelling, feed it image patches as if they were words, and change as little as possible.
What that removes is exactly the two biases above. Self-attention lets every patch look at every other patch from the very first layer, so there is no locality built in. There are no sliding filters, so there is no built-in translation equivariance either; position is something the model has to learn from a positional embedding. The only image-specific structure left is the cutting of the picture into a grid of patches at the start (a weak form of locality inside each patch) and the positional embeddings.
The consequence is one of the most important practical facts about ViTs. With less built-in knowledge, they need more data to learn what a CNN gets for free. The original paper found that ViTs trained on a mid-sized dataset like ImageNet-1k alone lagged behind comparable ResNets, and only overtook them when pre-trained on much larger datasets (Google's internal JFT-300M in that paper). Later work made ViTs much more data-efficient with better training recipes, augmentation and distillation, and we will meet DeiT below. But the lesson stands: a ViT trades inductive bias for flexibility, and pays for it in data.
For you as an engineer this translates to something concrete. You will rarely train a ViT from random weights. You will start from a strongly pre-trained backbone and fine-tune it or use it frozen. Keep that in mind; it shapes everything from model choice to edge deployment.
The ViT pipeline, step by step
Here is the whole forward pass in one paragraph before we take it apart. Cut the image into fixed-size square patches. Flatten each patch and project it to a vector, giving one token per patch. Prepend a special learnable token, [CLS]. Add a learned positional embedding to every token. Run the sequence through a stack of identical Transformer blocks, each with self-attention and a small MLP, wrapped in LayerNorm and residual connections. Take the final [CLS] vector and feed it to a linear layer that outputs class scores.
That is really all of it. Let us go through each part with numbers.
Patchify: an image becomes 196 tokens
Take the standard ViT-Base setup: a 224x224 RGB image and a patch size of 16x16.
patches per side = 224 / 16 = 14
number of patches = 14 x 14 = 196
values per patch = 16 x 16 x 3 = 768
So the image, a tensor of shape (3, 224, 224) with 150,528 numbers, becomes a table of 196 rows and 768 columns. That is 196 x 768 = 150,528 numbers again, of course: patchifying is a pure rearrangement, nothing is lost yet. Each row is one "word" of the image. The name "16x16 words" is literal.
The token count matters more than anything else in this article. Remember N = 196. Attention cost grows with the square of it.
Linear patch embedding: a strided convolution in disguise
Each flattened patch (768 numbers) is multiplied by a learned matrix of shape 768 x D to produce a D-dimensional token. For ViT-Base D = 768 as well, which is a coincidence of the design, not a requirement; the projection is 768 in, 768 out.
Now the trick you will see in every implementation. Flatten-then-Linear applied to non-overlapping patches is mathematically the same as a Conv2d with kernel_size=16 and stride=16 and 768 output channels. Why? A convolution with kernel 16 looks at a 16x16x3 window and computes a dot product with a filter of the same size. That is one output value per filter. With 768 filters you get 768 values per window. With stride 16 the windows do not overlap and tile the image exactly, which are precisely the patches. The filter weights, reshaped to 768 x 768, are the linear layer's weights.
In code you will write nn.Conv2d(3, 768, kernel_size=16, stride=16), and the output has shape (B, 768, 14, 14). Flatten the last two dimensions and swap axes to get (B, 196, 768). Our demo below proves the equivalence numerically, so you do not have to take my word for it.
The [CLS] token
We want one vector to represent the whole image so we can classify it. There are two common ways to get it: average all the output tokens, or add an extra token that exists only to collect information. The original ViT borrowed the second approach from BERT. It adds a learnable vector, [CLS], at the front of the sequence. It carries no image content, but because it takes part in attention in every layer, it can gather what it needs from the patch tokens. At the end, its output vector is fed to the classifier.
The sequence length is therefore 196 + 1 = 197 tokens. (Many modern backbones replace [CLS] with average pooling over patch tokens and work just as well. Either is fine, and pooling is often easier to export.)
Learned positional embeddings
Attention on its own has no idea where a token sits: shuffle the patches and, without positions, the output for each patch is the same. Yet for images, position matters. So we add to every token a learned vector that encodes its position. In ViT this is simply a parameter of shape (1, 197, 768), initialised with small random values and trained like any other weight. The model works out for itself that patch 17 sits below patch 3.
One practical consequence: those embeddings are tied to a 14x14 grid. If you fine-tune or run at a higher resolution, the number of patches changes, and you must resize the positional embeddings, usually by bicubic interpolation of the 14x14 grid to the new grid. Libraries such as timm do this for you. If you write your own, remember it, or your model will silently break at a new resolution.
Self-attention
This is the heart of the model. Every token is turned into three vectors by three learned linear maps: a query Q ("what am I looking for?"), a key K ("what do I contain?") and a value V ("what do I hand over if someone attends to me?"). Then:
Attention(Q, K, V) = softmax( Q Kα΅ / βd_k ) V
Read it left to right. Q is (N x d_k) and K is (N x d_k), so Q Kα΅ is an N x N table of scores: how strongly token i's query matches token j's key. Divide by βd_k, take a softmax along each row so the scores become weights that sum to 1, and use them to average the value vectors. Every output token is a weighted average of all the value vectors in the image, with weights chosen by content. That N x N table is where the "look at everything" flexibility comes from, and also where the cost comes from.
Why divide by βd_k
This is asked in interviews and is worth being able to explain. Suppose the components of a query q and a key k are independent, each with mean 0 and variance 1. Their dot product is a sum of d_k products:
q Β· k = qβkβ + qβkβ + ... + q_dk k_dk
each term: mean 0, variance 1 (var of a product of independent
zero-mean unit-variance values = 1)
sum of d_k: mean 0, variance d_k, standard deviation βd_k
So the raw scores get larger in magnitude as the head dimension grows. For d_k = 64 the typical score is around 8, and some are 20 or more. The softmax exponentiates its inputs, so with scores that spread out it collapses to almost a one-hot vector: one token gets a weight near 1 and the rest near 0. In that region the softmax gradient is almost zero, and learning stalls. Dividing by βd_k restores the variance to 1 whatever the head size, keeping the softmax in a range where gradients flow. We check this numerically in the NumPy section.
Multi-head attention
Instead of one big attention, ViT-Base splits the 768 dimensions into 12 heads of 64 dimensions each. Each head does the attention above on its own slice, and the results are concatenated and passed through one more linear layer. Different heads can specialise: one may follow edges and textures nearby, another may relate a wheel to the far-away car body. It costs about the same as one big head, and it gives the model several different "views" of the token relationships at once.
The MLP block
After attention, each token goes independently through a small two-layer network: expand from 768 to 3072 (a factor of 4), apply GELU, project back to 768. Attention mixes information between tokens; the MLP processes each token on its own. Most of a Transformer's parameters live in these MLPs, and as we will compute below, most of the compute does too.
LayerNorm, pre-norm and residuals
Each block wraps both sub-layers in a residual connection, and applies LayerNorm before each sub-layer rather than after. In ViT the block reads:
x = x + Attention( LayerNorm(x) )
x = x + MLP( LayerNorm(x) )
The residual (x + ...) gives gradients a clean path back through 12 or more layers, so deep stacks train. LayerNorm normalises each token's vector across its own features, independent of the batch, which suits sequence models and small-batch inference. Placing it before the sub-layer ("pre-norm") makes training more stable than the original post-norm arrangement, which is why nearly every modern Transformer does it. A final LayerNorm is applied after the last block.
The classification head
Take the output at position 0 (the [CLS] token), and apply one Linear(768, num_classes). Softmax of those logits gives class probabilities. For fine-tuning you replace only this layer with one sized for your classes, and often nothing else.
Attention in NumPy, no framework
Before the PyTorch version, here is single-head attention using only NumPy. If you can read this, you can read any attention implementation.
import numpy as np
def softmax(x, axis=-1):
x = x - x.max(axis=axis, keepdims=True) # subtract the max for numerical safety
e = np.exp(x)
return e / e.sum(axis=axis, keepdims=True)
def single_head_attention(x, Wq, Wk, Wv):
"""x: (N, D) tokens. Wq/Wk/Wv: (D, d_k). Returns (N, d_k) and the (N, N) weights."""
Q, K, V = x @ Wq, x @ Wk, x @ Wv
d_k = Q.shape[-1]
scores = Q @ K.T / np.sqrt(d_k) # (N, N): how much token i looks at token j
weights = softmax(scores, axis=-1) # each row sums to 1
return weights @ V, weights # weighted average of value vectors
rng = np.random.default_rng(0)
N, D, d_k = 5, 8, 4
x = rng.normal(size=(N, D))
Wq, Wk, Wv = (rng.normal(size=(D, d_k)) / np.sqrt(D) for _ in range(3))
out, w = single_head_attention(x, Wq, Wk, Wv)
print("output shape:", out.shape)
print("weights shape:", w.shape)
print("row sums:", w.sum(axis=-1).round(3))
# Why the sqrt(d_k) scale: raw dot products of unit-variance vectors
for d in (16, 64, 256, 1024):
q = rng.normal(size=(20000, d)); k = rng.normal(size=(20000, d))
dots = (q * k).sum(axis=1)
print(d, round(dots.std(), 2), round((dots / np.sqrt(d)).std(), 2))
Running it prints, as I verified while writing this article:
output shape: (5, 4)
weights shape: (5, 5)
row sums: [1. 1. 1. 1. 1.]
16 4.01 1.0
64 8.01 1.0
256 15.97 1.0
1024 31.97 1.0
The last four lines are the variance argument made visible. The standard deviation of the raw dot product is βd (4, 8, 16, 32), and after the scale it is 1 every time. Also notice the x - x.max(...) line inside softmax: it does not change the result, but it stops exp from overflowing, a trick you will meet again when quantising.
Multi-head attention is this same function run on several slices of the features. Time to build the real thing.
A complete ViT in PyTorch
This is a full, working ViT. It is written to be read, not to win benchmarks. I tested it on PyTorch 2.7; nothing here needs a particular version.
import torch
import torch.nn as nn
import torch.nn.functional as F
class PatchEmbed(nn.Module):
"""Cut the image into patches and project each one to `dim` numbers."""
def __init__(self, img_size=224, patch=16, in_ch=3, dim=768):
super().__init__()
self.grid = img_size // patch
self.num_patches = self.grid ** 2
# A conv with kernel = stride = patch is "flatten each patch, then Linear".
self.proj = nn.Conv2d(in_ch, dim, kernel_size=patch, stride=patch)
def forward(self, x): # (B, 3, H, W)
x = self.proj(x) # (B, dim, H/p, W/p)
return x.flatten(2).transpose(1, 2) # (B, N, dim)
class Attention(nn.Module):
def __init__(self, dim, heads):
super().__init__()
assert dim % heads == 0
self.heads, self.dh = heads, dim // heads
self.qkv = nn.Linear(dim, dim * 3)
self.proj = nn.Linear(dim, dim)
def forward(self, x): # (B, N, D)
B, N, D = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dh)
q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) # each (B, heads, N, dh)
att = (q @ k.transpose(-2, -1)) / self.dh ** 0.5 # (B, heads, N, N)
att = att.softmax(dim=-1)
out = (att @ v).transpose(1, 2).reshape(B, N, D) # merge heads
return self.proj(out)
class Block(nn.Module):
def __init__(self, dim, heads, mlp_ratio=4.0):
super().__init__()
self.norm1 = nn.LayerNorm(dim, eps=1e-6)
self.attn = Attention(dim, heads)
self.norm2 = nn.LayerNorm(dim, eps=1e-6)
hidden = int(dim * mlp_ratio)
self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Linear(hidden, dim))
def forward(self, x):
x = x + self.attn(self.norm1(x)) # pre-norm + residual
x = x + self.mlp(self.norm2(x))
return x
class ViT(nn.Module):
def __init__(self, img_size=224, patch=16, num_classes=1000,
dim=768, depth=12, heads=12, mlp_ratio=4.0):
super().__init__()
self.patch_embed = PatchEmbed(img_size, patch, 3, dim)
n = self.patch_embed.num_patches
self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
self.pos_embed = nn.Parameter(torch.zeros(1, n + 1, dim))
nn.init.trunc_normal_(self.pos_embed, std=0.02)
nn.init.trunc_normal_(self.cls_token, std=0.02)
self.blocks = nn.Sequential(*[Block(dim, heads, mlp_ratio) for _ in range(depth)])
self.norm = nn.LayerNorm(dim, eps=1e-6)
self.head = nn.Linear(dim, num_classes)
def forward(self, x):
x = self.patch_embed(x) # (B, N, D)
cls = self.cls_token.expand(x.shape[0], -1, -1) # (B, 1, D)
x = torch.cat([cls, x], dim=1) + self.pos_embed # (B, N+1, D)
x = self.blocks(x)
x = self.norm(x)
return self.head(x[:, 0]) # classify from [CLS]
Walk through the shapes in Attention, because that is where people get lost. One Linear(dim, 3*dim) produces Q, K and V in one matrix multiply. We reshape to (B, N, 3, heads, dh), permute so the "3" comes first, and unbind it into three tensors of shape (B, heads, N, dh). Now q @ k.transpose(-2, -1) is a batched matrix multiply of (N x dh) by (dh x N), giving (B, heads, N, N): one attention table per head per image. After the softmax, multiplying by V gives (B, heads, N, dh); we move the heads next to the feature axis and reshape back to (B, N, D) to concatenate them, then apply the output projection.
I used unbind(0) rather than tuple-unpacking the tensor on purpose. Unpacking a tensor works in eager mode, but tracing-based exporters (like the ONNX exporter we use later) warn about it because it looks like a Python loop over a tensor. unbind traces cleanly.
Now a demo that prints the shapes, proves the convolution claim, and trains a tiny ViT to check that gradients flow:
if __name__ == "__main__":
torch.manual_seed(0)
model = ViT()
print("params:", sum(p.numel() for p in model.parameters()))
x = torch.randn(2, 3, 224, 224)
pe = model.patch_embed(x); print("patch tokens:", tuple(pe.shape))
logits = model(x); print("logits:", tuple(logits.shape))
# Patch embedding really is "flatten each patch, then Linear"
p = 16
patches = x.unfold(2, p, p).unfold(3, p, p) # (B, 3, 14, 14, 16, 16)
patches = patches.permute(0, 2, 3, 1, 4, 5).reshape(2, 196, 3 * p * p)
W = model.patch_embed.proj.weight.reshape(768, -1) # (768, 768)
lin = patches @ W.T + model.patch_embed.proj.bias
print("conv == linear:", torch.allclose(lin, pe, atol=1e-4))
# A tiny ViT can memorise 64 random images: checks the whole thing trains
tiny = ViT(img_size=32, patch=8, num_classes=4, dim=64, depth=2, heads=4)
opt = torch.optim.AdamW(tiny.parameters(), lr=1e-3)
xs = torch.randn(64, 3, 32, 32); ys = torch.randint(0, 4, (64,))
for step in range(200):
loss = F.cross_entropy(tiny(xs), ys)
opt.zero_grad(); loss.backward(); opt.step()
if step % 50 == 0 or step == 199:
print(step, round(loss.item(), 3))
# Our attention should match PyTorch's fused implementation
q, k, v = (torch.randn(1, 4, 10, 16) for _ in range(3))
a = ((q @ k.transpose(-2, -1)) / 16 ** 0.5).softmax(-1) @ v
b = F.scaled_dot_product_attention(q, k, v)
print("sdpa match:", torch.allclose(a, b, atol=1e-5))
Output when I ran it:
params: 86567656
patch tokens: (2, 196, 768)
logits: (2, 1000)
conv == linear: True
0 1.516
50 0.007
100 0.004
150 0.003
199 0.002
sdpa match: True
Three things to notice. The parameter count, 86,567,656, is the well-known figure for ViT-B/16, and we are about to derive it by hand. The patch tokens are (2, 196, 768), exactly the arithmetic from earlier. And the loss on random labels falling to nearly zero is memorisation, not intelligence; it just proves that every part of the graph receives gradient. (The exact loss values depend on your seed and PyTorch build; treat them as a sanity check, not a target.)
The last check matters for production. F.scaled_dot_product_attention computes exactly our formula, but PyTorch can dispatch it to a fused, memory-efficient kernel (FlashAttention-style) that never builds the full N x N table in memory. In real code you should call it instead of the three manual lines. I wrote it out by hand above so that the math was visible.
Counting parameters and FLOPs by hand
Estimating cost before you run anything is the most useful habit in edge work. Let us do ViT-Base: D = 768, 12 layers, 12 heads, MLP width 3072, 224x224 input, 16x16 patches, 1000 classes.
Parameters. A linear layer from a to b has a x b weights plus b biases.
Patch embedding : 768 x 768 + 768 = 590,592
[CLS] token : 768 = 768
Position embed : 197 x 768 = 151,296
One block:
LayerNorm x2 : 2 x (2 x 768) = 3,072
QKV linear : 768 x 2304 + 2304 = 1,771,776
Output proj : 768 x 768 + 768 = 590,592
MLP up : 768 x 3072 + 3072 = 2,362,368
MLP down : 3072 x 768 + 768 = 2,360,064
Block total = 7,087,872
12 blocks = 85,054,464
Final LayerNorm = 1,536
Head : 768 x 1000 + 1000 = 769,000
Total = 86,567,656
That matches the count our code printed, which is a satisfying check that the mental model and the implementation agree. A handy rule: each block has roughly 12 DΒ² weights (3DΒ² for QKV, DΒ² for the projection, 8DΒ² for the MLP), so 12 blocks hold about 12 x 12 x 768Β² β 85 million.
FLOPs. Count multiply-accumulates (MACs); one MAC is two floating-point operations. A linear layer applied to N tokens costs N x a x b MACs. Per block, the weight matrices contribute about 12 DΒ² MACs per token:
Linear layers per block : 12 x 768Β² x 197 β 1.394 billion MACs
Attention matmuls : Q Kα΅ = NΒ² x D = 197Β² x 768 β 29.8 million
A V = NΒ² x D β 29.8 million
sum β 59.6 million
One block β 1.454 billion MACs
12 blocks β 17.45 billion MACs
Patch embedding : 196 x 768 x 768 β 0.116 billion MACs
Total β 17.6 billion MACs β 35 GFLOPs (approximate)
This lines up with the roughly 17.6 GMACs commonly quoted for ViT-B/16. I am ignoring softmax, LayerNorm, GELU and the residual adds, which are small in FLOPs but not necessarily in time, as we will see on the device. Treat 35 GFLOPs per 224x224 image as approximate.
Two readings of this arithmetic are worth keeping. First, at 224x224 the attention matrix products are only about 4% of the compute (59.6 / 1454). The linear layers dominate, and the MLPs alone are two thirds of those. So at normal resolution, "attention is expensive" is not the right intuition. Second, watch what happens when N grows.
The attention cost problem
The NΒ² in NΒ² x D comes from the N x N score table, and it changes the picture when images get bigger. The linear-layer cost grows in proportion to N; the attention product grows with NΒ². Setting them equal, 2 NΒ² D = 12 DΒ² N, gives N = 6D. For D = 768 the crossover is at N = 4,608 tokens: beyond that, attention matrix products cost more than all the linear layers combined.
Double the resolution of the image and you get four times the tokens, and sixteen times the size of the attention table. Some concrete cases with the same patch size of 16:
Image size Tokens (N) Attention table per head vs 224x224
224 x 224 196 ~38,400 entries 1x
384 x 384 576 ~331,800 entries ~8.6x
1024 x 1024 4,096 ~16,800,000 entries ~440x
(Add one for [CLS]; the table above ignores it for readability.) At 1024x1024, with ViT-Base's 12 heads in FP16, a naive implementation that materialises the table needs 12 x 4096Β² x 2 bytes, about 400 MB, for one layer's attention weights, and it has to be written to memory and read back. On an edge device, where memory bandwidth is the scarce resource, that is what hurts. Fused attention kernels avoid storing the table, which fixes the memory but not the arithmetic.
This matters for anything that must see detail: a farm drone photographing crops from 30 metres, a factory camera hunting for a hairline scratch, a warehouse camera reading small print. You cannot always just shrink the image; the defect disappears. Here is how the field responded.
Shift the problem: windows (Swin)
Swin Transformer (arXiv:2103.14030) computes attention only inside small local windows of patches (7x7 in the standard configuration), so cost grows linearly with the number of tokens instead of quadratically. To let information cross window borders, alternate layers shift the window grid. It also builds a hierarchy: patches are merged as you go deeper, so early layers work at fine resolution and later ones at coarse resolution, like a CNN feature pyramid. That makes Swin a natural backbone for detection and segmentation. Notice that it does so by putting locality back: an admission that the CNN's bias was useful.
Shift the problem: get data efficiency by teaching
DeiT (arXiv:2012.12877) showed that a ViT trained on ImageNet-1k alone can be competitive, given strong augmentation and regularisation plus distillation: a distillation token learns to match a teacher network's predictions, typically a CNN. This is where the "smaller model" rung of the edge ladder starts: DeiT-Tiny and DeiT-Small are far cheaper than ViT-Base and are common starting points for deployment.
Shift the problem: hybrid conv and transformer models
If the CNN's locality is cheap and good at early layers, keep it there. Hybrid designs use convolutions to downsample the image quickly and run attention only on a small number of tokens. MobileViT (arXiv:2110.02178) mixes MobileNet-style convolution blocks with small Transformer blocks that operate on patches within feature maps, aiming at mobile hardware. Others in the same family (EfficientViT, and Transformer-flavoured MobileNet variants) follow the same principle. The general recipe: convolutions where they are cheap and effective, attention where global context pays off. On many edge accelerators, which were designed with convolutions in mind, these hybrids can be easier to run efficiently than a pure ViT, though you should measure on your own device rather than trust a paper's table.
Shift the problem: reuse a frozen backbone
Often the best way to avoid training a large ViT is not to train it at all. DINOv2 (arXiv:2304.07193) is a family of ViTs (with patch size 14) trained by self-supervision on a large curated image set, producing general-purpose features that work well when frozen: you train only a small linear or MLP head, or use nearest-neighbour search on the embeddings. For a factory that has 200 labelled defect images, this is a far better plan than fine-tuning a big model. It also gives you a clean split for edge deployment: one shared frozen backbone on the device, with tiny task-specific heads that you can swap without touching the heavy part.
A brief word on vision-language models
CLIP (arXiv:2103.00020) trains an image encoder (often a ViT) and a text encoder together so that a picture and its caption land near each other in a shared embedding space; SigLIP (arXiv:2303.15343) is a variant with a sigmoid-based loss. The practical payoff is zero-shot recognition: describe a class in words ("a crushed cardboard box") and compare its text embedding with the image embedding. Those image encoders are ViTs, so everything in this article about exporting and optimising them applies, and they are handy for bootstrapping a dataset. They are not a replacement for a supervised model where you need calibrated, high-precision decisions.
Compute smarter, not only smaller
Two more levers deserve a mention. FlashAttention-style kernels (arXiv:2205.14135) reorganise the attention computation to avoid writing the N x N table to slow memory; you get them by calling F.scaled_dot_product_attention, and by using a runtime that fuses attention (TensorRT does for supported patterns). And token reduction methods drop or merge unimportant patch tokens as the network gets deeper, so later layers see fewer tokens. These are active research areas; check that a technique has a well-supported export path before you build on it.
Physical AI: what changes when the model leaves the notebook
"Physical AI" is the name for AI that acts in the world through a body: robots, drones, vehicles, cameras that trigger machinery. The models are the same ones we have been discussing. The environment is not, and it changes the engineering rules.
Latency is a hard budget, not a metric. A camera at 30 frames per second gives you 33 ms per frame for everything: capture, resize, inference, post-processing, and the decision. A drone moving at 10 m/s travels a metre in 100 ms, so a model that takes 100 ms is looking at where the drone was. What matters is not the average but the slowest frames, because a controller reacts to those. This is why we measure the 95th percentile and beyond, not the mean.
Power and heat are limits. A board that draws 15 watts in a sealed enclosure in the sun will throttle its clocks when it gets hot. A benchmark that runs for ten seconds on a cool desk tells you almost nothing about performance after twenty minutes in a warehouse. Battery-powered devices also convert every millijoule into flight time or driving range.
Memory is small and its bandwidth is the real ceiling. Many edge devices share one pool of memory between the CPU and GPU. Batch size 1 is normal (you are processing the frame that just arrived). The weights of ViT-Base are about 346 MB in FP32, 173 MB in FP16 and 87 MB in INT8, and each inference must stream them through the chip. As an illustration, if your device could sustain 50 GB/s of memory traffic, reading 173 MB of FP16 weights would take at least 3.5 ms even with infinitely fast arithmetic. Shrinking the weights and the activations speeds things up because it moves fewer bytes, not only because integer maths is faster.
Small operations cost real time. At batch 1, the big matrix multiplies are reasonably efficient, but LayerNorm, softmax, GELU and reshapes are memory-bound: each one reads and writes a whole tensor for very little arithmetic. Running them as separate steps is wasteful. This is a large part of why a compiler like TensorRT, which fuses these into fewer kernels, gives such big speedups.
There may be no network. A robot in a cold store, a drone behind a hill, or a camera in a plant with no outbound connection cannot call an API. The model, its fallback behaviour and its monitoring must all live on the device. You also cannot fix a bad model by redeploying a server; updates are shipped to a fleet.
Mistakes have physical consequences. A wrong label on a web page is an annoyance. A wrong label on a robot means the arm drops something, or the drone flies into a field boundary. So the design question changes from "how accurate is it?" to "what does it do when it is wrong or unsure?" We come back to that in the production section.
The optimisation ladder
You do not need every technique at once. Climb one rung at a time, measuring accuracy and latency at each step, and stop as soon as you meet your budget. The order below is roughly from cheapest and safest to hardest and riskiest.
- Choose a smaller model or a lower resolution. The biggest win is the one that costs no engineering.
- Export to ONNX so the model is a portable graph rather than Python code.
- Run it in a fast runtime (ONNX Runtime, then TensorRT on NVIDIA hardware).
- Lower precision: FP16 first, then INT8 with calibration.
- Prune or distil if you still need more, and you have training budget.
Rung 1: a smaller model
ViT-Base at 35 GFLOPs is a workstation model. Consider the ladder of options: DeiT-Tiny or Small, a small DINOv2 backbone (ViT-S/14), a hybrid such as MobileViT, or even a modern CNN. Parameters are roughly proportional to depth x DΒ², and per-image compute to that times the token count. Halving D quarters most of the cost. Reducing the input from 224 to 160 pixels with patch size 16 gives 100 tokens instead of 196, close to half the linear-layer compute and far less attention. Whether the accuracy loss is acceptable is an empirical question about your task, not something to guess. If your objects fill much of the frame, a lower resolution may cost you nothing; if your defects are tiny, it will cost you everything.
Do not skip a CNN baseline. A small, well-tuned CNN is sometimes the right answer for a fixed camera in a controlled scene, and having its numbers tells you what the ViT is buying you.
Rung 2: export to ONNX
ONNX is a framework-neutral file format for a computation graph. Once your model is in that format, many runtimes and compilers can consume it. You need the onnx package installed (pip install onnx onnxruntime). Newer PyTorch releases default torch.onnx.export to the torch.export-based exporter, which also needs onnxscript; older releases used a TorchScript tracing exporter. Passing the dynamo flag explicitly keeps your script's behaviour clear across versions. I tested the following with dynamo=False on PyTorch 2.7:
import torch
from vit import ViT # the model defined above, saved as vit.py
model = ViT(img_size=224, patch=16, num_classes=10, dim=192, depth=4, heads=3).eval()
dummy = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model, dummy, "vit_tiny.onnx",
input_names=["image"], output_names=["logits"],
opset_version=17,
dynamo=False,
dynamic_axes={"image": {0: "batch"}, "logits": {0: "batch"}},
)
Two habits to build now. Always call .eval() before exporting, so that dropout and similar training-only behaviour is off. And always verify the exported graph against the original on real inputs, because an export that runs is not the same as an export that is correct:
import numpy as np
import onnxruntime as ort
sess = ort.InferenceSession("vit_tiny.onnx", providers=["CPUExecutionProvider"])
x = np.random.randn(3, 3, 224, 224).astype(np.float32)
onnx_out = sess.run(None, {"image": x})[0]
with torch.no_grad():
ref = model(torch.from_numpy(x)).numpy()
print("max abs diff:", np.abs(onnx_out - ref).max())
For me this printed a difference around 1e-6, which is float rounding. Anything much larger than about 1e-4 in FP32 deserves an investigation before you go any further.
The most common export problems are unsupported operators (usually custom or very new PyTorch ops), Python control flow that depends on tensor values (the tracer freezes one branch), and shapes that were silently frozen to the dummy input's size. If your model uses a Python if on a tensor, or .item(), expect trouble. Keep the model's forward path plain tensor operations and it will export cleanly. This is another reason to use F.scaled_dot_product_attention or plain matmul attention rather than exotic custom kernels: the ops are standard.
Rung 3: a fast runtime
ONNX Runtime (onnxruntime.ai) loads the file and applies graph optimisations such as constant folding and operator fusion, then executes it through an "execution provider" for your hardware: CPU, CUDA, TensorRT, CoreML, DirectML, or NNAPI on Android, among others. The providers list is ordered by preference and falls back down the list:
sess = ort.InferenceSession(
"vit_tiny.onnx",
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
print(sess.get_providers()) # check which one was actually used
Print that last line every time. If the CUDA provider fails to load (a common cause is a mismatched CUDA or cuDNN version), ONNX Runtime falls back to the CPU quietly, and you will spend an afternoon wondering why your GPU is idle.
Rung 4a: how quantisation works
Quantisation stores numbers with fewer bits. Instead of a 32-bit float per weight, you store an 8-bit integer plus a shared scale that lets you approximately recover the float. Weights take 4x less memory than FP32, the memory traffic drops by the same factor, and on hardware with integer matrix units the arithmetic is faster too.
The standard scheme is affine (asymmetric) quantisation. You pick a scale s (a positive float) and a zero-point z (an integer), and map real values x to integers q with:
q = clip( round(x / s) + z, q_min, q_max )
xΜ = s * (q - z) (dequantised approximation of x)
for unsigned 8-bit: q_min = 0, q_max = 255
s = (x_max - x_min) / 255
z = round( q_min - x_min / s )
The zero-point exists so that the real value 0.0 maps to an exact integer. That matters because zero padding and ReLU outputs are everywhere. Symmetric quantisation is the special case with z = 0 and a range centred on zero; it is simpler and often used for weights.
A worked example. Suppose a tensor has values between -1.0 and 2.0. Then s = 3.0 / 255 β 0.01176 and z = round(0 - (-1.0)/0.01176) = round(85.0) = 85. Quantise five values:
x : -1.0 -0.2 0.0 0.5 2.0
q : 0 68 85 127 255
xΜ : -1.0 -0.2 0.0 0.4941 2.0
The value 0.5 comes back as 0.4941: an error of about 0.006, which is at most half a step (s / 2 β 0.0059). That bounded rounding error is the price you pay. Here is the same idea in NumPy, including what happens when a single outlier stretches the range:
import numpy as np
def quant_params(x_min, x_max, n_bits=8):
x_min, x_max = min(x_min, 0.0), max(x_max, 0.0) # the range must contain 0
q_min, q_max = 0, 2**n_bits - 1
scale = (x_max - x_min) / (q_max - q_min)
zero_point = int(round(q_min - x_min / scale))
return scale, int(np.clip(zero_point, q_min, q_max))
def quantise(x, scale, zp, n_bits=8):
return np.clip(np.round(x / scale) + zp, 0, 2**n_bits - 1).astype(np.uint8)
def dequantise(q, scale, zp):
return scale * (q.astype(np.float32) - zp)
rng = np.random.default_rng(0)
x = rng.normal(0.0, 1.0, size=10_000).astype(np.float32)
x[0] = 40.0 # a single outlier
s, z = quant_params(x.min(), x.max()) # naive: use the full min/max range
xh = dequantise(quantise(x, s, z), s, z)
lo, hi = np.percentile(x, [0.1, 99.9]) # calibrated: ignore the extreme tails
s2, z2 = quant_params(lo, hi)
xh2 = dequantise(quantise(x, s2, z2), s2, z2)
bulk = np.abs(x) < 3 # the values that carry most of the signal
print("step size naive %.4f clipped %.4f" % (s, s2))
print("bulk error naive %.4f clipped %.4f" % (np.abs(x - xh)[bulk].mean(),
np.abs(x - xh2)[bulk].mean()))
When I ran it, the step size went from 0.1722 with the naive range to 0.0240 with the clipped range, and the average error on ordinary values fell from about 0.043 to about 0.006. One outlier at 40 made every other value 7x coarser. That is the whole story of calibration.
Rung 4b: calibration and why ViTs are fussy
Weights are known ahead of time, so their ranges are easy to measure. Activations (the tensors flowing through the network) depend on the input, so you cannot know their ranges until you see data. Post-training static quantisation solves this by running a small calibration set (typically a few hundred representative images) through the model, recording each activation's range or histogram, and choosing the scale and zero-point from it, often by clipping the tails (percentile or entropy-based methods) rather than using the raw min and max.
Two rules follow. The calibration set must come from your deployment domain: images from the actual camera, in the actual lighting. Calibrating on stock photos and running in a dim warehouse is how you get an accuracy collapse that no unit test catches. And you need a proper held-out evaluation set, separate from calibration, to measure the accuracy drop.
Transformers are harder to quantise than CNNs. Their activations, particularly after LayerNorm and in the MLPs, contain large outliers in a few channels, which is exactly the "one big value stretches the range" problem above. Softmax outputs sit in [0, 1] with a very skewed distribution. So the usual practice is to quantise the heavy matrix multiplies (the QKV, projection and MLP linear layers, which is where nearly all the FLOPs are) and leave LayerNorm and softmax in higher precision, or to use per-channel scales for weights. If accuracy drops more than you can tolerate, find the layers that are the most sensitive by quantising one at a time, and keep those in FP16. The background paper for this style of integer inference is Jacob et al. (arXiv:1712.05877).
ONNX Runtime has a static quantisation tool. It calls a data reader for calibration batches:
import numpy as np
from onnxruntime.quantization import (
quantize_static, CalibrationDataReader,
QuantFormat, QuantType, CalibrationMethod,
)
class ImageReader(CalibrationDataReader):
def __init__(self, input_name, batches):
self._it = iter([{input_name: b} for b in batches])
def get_next(self):
return next(self._it, None) # None tells the tool "no more data"
# In real use: 100-500 preprocessed images from the real camera, each (1, 3, 224, 224) float32.
calib = [np.random.randn(1, 3, 224, 224).astype(np.float32) for _ in range(32)]
quantize_static(
"vit_tiny.onnx", "vit_tiny.int8.onnx",
ImageReader("image", calib),
quant_format=QuantFormat.QDQ, # explicit Quantize/Dequantize nodes; works well with TensorRT
activation_type=QuantType.QInt8,
weight_type=QuantType.QInt8,
per_channel=True, # one scale per output channel for weights
calibrate_method=CalibrationMethod.MinMax, # Entropy and Percentile also exist
)
I ran this on the tiny model and it worked, producing a file about a quarter of the FP32 file's size. The random calibration data above is only there to make the snippet self-contained: it produces a meaningless calibration and you must never ship a model calibrated on noise. Also expect a warning suggesting a pre-processing step; ONNX Runtime ships a pre-process tool (python -m onnxruntime.quantization.preprocess) that runs shape inference and optimisation before quantising, and it is worth using on real models.
Treat quantisation as a hypothesis you test: quantise, then evaluate accuracy on your held-out set, and compare per-class results, not just the overall number. A drop from 96.0% to 95.5% overall can hide a class that went from 90% to 60%. If the drop is too large, options include a better calibration set, a different calibration method, skipping sensitive layers, or quantisation-aware training, where you fine-tune the model with simulated quantisation in the loop so it learns to be robust to it.
Rung 4c: FP16 first
Half precision is the easy win on GPUs. Weights and activations take half the bytes, and modern NVIDIA GPUs have tensor cores built for FP16 matrix maths. The dynamic range is much smaller than FP32 (the largest FP16 value is 65,504), so overflow is possible in extreme activations, but for most vision transformers it works with little or no accuracy loss. Try FP16 before INT8: it is a one-flag change in TensorRT and needs no calibration data. Some devices also support BF16, which keeps FP32's range with lower precision.
Rung 5: TensorRT on Jetson-class hardware
TensorRT (developer.nvidia.com/tensorrt) is NVIDIA's inference compiler. It reads an ONNX graph, fuses operations, chooses the fastest kernel for each layer on your specific GPU, and can use FP16 and INT8 tensor cores. On NVIDIA's Jetson boards (the Orin series and newer), it is the standard path. Its command line tool trtexec builds an engine and benchmarks it in one go:
# Build an FP16 engine on the target device
trtexec --onnx=vit_tiny.onnx \
--saveEngine=vit_tiny_fp16.plan \
--fp16 \
--shapes=image:1x3x224x224
# INT8 needs calibration data or an ONNX file that already carries Q/DQ nodes
trtexec --onnx=vit_tiny.int8.onnx \
--saveEngine=vit_tiny_int8.plan \
--int8 --fp16 \
--shapes=image:1x3x224x224
trtexec prints latency statistics including the median and high percentiles, which is handy, but it measures only the inference step, not your whole pipeline. A few things to know before you rely on it:
- An engine is tied to the GPU and the TensorRT version that built it. Build on the device you deploy to (or an identical one), and rebuild when you upgrade JetPack or TensorRT. Ship the ONNX file as the portable artefact and the engine as a cache.
- Power mode matters on Jetson. The
nvpmodeltool selects a power profile andjetson_clockspins the clocks at their maximum. Decide which mode you will use in the field and benchmark in that mode. Numbers from a maximum-performance mode on a bench power supply may be unreachable on the battery in the robot.tegrastatsshows live utilisation, memory, temperature and power. - The first run is slow. TensorRT and CUDA do a lot of setup lazily. This is why our benchmark has a warm-up phase.
- Fused attention may need supported patterns. Recent TensorRT versions fuse common attention patterns. If your graph uses an unusual formulation, it may not fuse, and it will be slower. Compare against a standard formulation.
For phones and other targets the story is similar with different tools. TensorFlow Lite (now maintained under the LiteRT name) and Core ML on Apple devices each have their own converters, their own supported-operator lists and their own quantisation flows. The general lesson holds: export a clean graph, test the operators, and keep the converted model as a build artefact you can regenerate, never a hand-edited file.
Rung 6: pruning and distillation, briefly
Pruning removes weights or whole structures. Unstructured pruning zeroes out individual weights; it can reach high sparsity but usually gives no speedup on ordinary hardware, because dense kernels do the same work regardless of the zeros. Structured pruning removes entire attention heads, MLP neurons or layers, which really does shrink the compute, but needs retraining to recover accuracy. For ViTs, removing heads or reducing MLP width are the practical choices.
Distillation trains a small "student" to imitate a large "teacher": rather than learning only from hard labels, it learns from the teacher's full probability distribution, which carries much more information ("this is 90% a crushed box, 8% an open box"). When you have a strong large model and a tight latency budget, distilling into a DeiT-Small-size or hybrid student is often the highest-value step after the free ones. You also get to use unlabelled images, as the teacher supplies the labels.
Both cost training time, which is why they sit at the top of the ladder.
A latency benchmark that does not lie
Most bad latency numbers come from bad benchmarks. The classic mistakes are timing a single run, including one-time setup, forgetting that GPU work is asynchronous, and reporting the mean. On a GPU, model(x) returns immediately after queuing the work; if you stop the timer at that point you measure the launch, not the computation. You must synchronise before reading the clock.
This harness handles the four essentials: warm-up runs, synchronisation, many iterations, and percentile reporting.
import time
import statistics
import torch
def benchmark(fn, warmup=30, iters=300, use_cuda=True):
"""Time fn() and return (median_ms, p95_ms, max_ms).
fn must run one complete inference, including any device copies you
want counted. Pass a closure so the same harness works for PyTorch,
ONNX Runtime and TensorRT.
"""
sync = torch.cuda.synchronize if (use_cuda and torch.cuda.is_available()) else (lambda: None)
for _ in range(warmup): # lazy init, autotuning, cache warm-up
fn()
sync()
times = []
for _ in range(iters):
sync() # make sure earlier work has finished
t0 = time.perf_counter()
fn()
sync() # wait for THIS call's GPU work
times.append((time.perf_counter() - t0) * 1000.0)
times.sort()
p95 = times[int(0.95 * len(times)) - 1]
return statistics.median(times), p95, times[-1]
# Example with PyTorch on a GPU (FP16, batch size 1):
if torch.cuda.is_available():
model = ViT().eval().cuda().half()
x = torch.randn(1, 3, 224, 224, device="cuda", dtype=torch.float16)
@torch.inference_mode()
def run():
return model(x)
med, p95, worst = benchmark(run)
print(f"median {med:.2f} ms | p95 {p95:.2f} ms | worst {worst:.2f} ms")
I ran a CPU-only variant of this harness against a small ONNX Runtime model and it returned sensible values; I am not going to print numbers for a device you do not own. Run it on yours. Some notes on interpreting the result:
- Report the median and p95 (and the worst case), not the mean. One slow outlier can distort a mean, and the tail is what causes dropped frames. If p95 is twice the median, find out why (thermal throttling, memory allocation, other processes) before you tune anything else.
- Benchmark long enough to heat the device. Run a five-to-ten-minute soak test in the enclosure you will actually use, log the temperature and clock speed, and compare the first minute against the last. The gap is your throttling penalty.
- Use realistic inputs and the exact runtime. Random tensors are fine for timing dense maths, but the full pipeline includes camera capture, colour conversion, resizing and post-processing, and those often take longer than the model. Time them too.
- For ONNX Runtime on GPU, the ordinary
session.runcopies inputs from host memory and outputs back. If you want to measure the model alone, use I/O binding with device tensors; if you want the honest end-to-end figure, include the copies. - Measure the accuracy alongside the latency, always. A benchmark table with speeds and no accuracy column is how a broken quantised model gets shipped.
An illustrative shape for the final table you shoudl aim to produce (these are placeholders for your own measurements, not results):
Variant Size (MB) Median (ms) p95 (ms) Top-1 on YOUR test set
FP32 PyTorch ? ? ? ?
FP16 TensorRT ? ? ? ?
INT8 TensorRT ? ? ? ?
Fill in every cell from your device or do not use the table.
Production reality
A model that scores well on a test set is the start of the job. In the field, three things happen: the world drifts away from the training data, the pipeline around the model turns out to matter as much as the model, and you find yourself needing to update hundreds of devices safely.
Distribution shift
Your training data was a snapshot. The warehouse changes its lighting from fluorescent to LED. A new camera model has a different sensor and lens. Winter arrives, and low sun casts long shadows across the field the drone maps. Dust builds up on the lens. Suppliers change their box design. The model's accuracy quietly decays, and there is no exception to catch: it produces confident, wrong answers.
The defence is to make drift visible. You usually cannot compute accuracy in the field, because you have no labels, but you can watch things you can compute: input statistics (mean brightness, contrast, a sharpness measure such as the variance of the Laplacian, which drops with motion blur or a dirty lens), and output statistics (the distribution of predicted classes, the average confidence, the fraction of low-confidence frames). Record reference values from your training or validation data, keep rolling windows on the device, and raise an alarm when the window drifts too far from the reference. A sudden drop in mean confidence on a device that was fine yesterday is more often a lens problem than a model problem, and it is cheap to detect.
Calibration: is 90% really 90%?
A network's softmax output looks like a probability but usually is not one. Modern networks tend to be overconfident: when they say 95%, they may be right only 85% of the time (Guo et al., "On Calibration of Modern Neural Networks", arXiv:1706.04599). If your robot's behaviour depends on a threshold, that threshold only means something if the confidence is calibrated.
The simplest fix is temperature scaling: divide the logits by a single learned number T before the softmax, chosen to minimise the log-loss on a held-out set. It does not change which class wins, only how confident it claims to be. Refit it after any quantisation step, since quantisation changes the logits' scale.
import torch
import torch.nn.functional as F
def fit_temperature(logits, labels):
"""Learn one number T so softmax(logits / T) is honest. Fit on held-out data."""
log_t = torch.zeros(1, requires_grad=True) # optimise log T so T stays positive
opt = torch.optim.LBFGS([log_t], lr=0.1, max_iter=100)
def closure():
opt.zero_grad()
loss = F.cross_entropy(logits / log_t.exp(), labels)
loss.backward()
return loss
opt.step(closure)
return log_t.exp().item()
Evaluate calibration with a reliability diagram, or a summary such as expected calibration error, before and after. And remember the limits: temperature scaling fixes the confidence on data like your validation set. It cannot make the model know that it is looking at something it has never seen. A model can be perfectly calibrated on its test set and still be 99% sure that a photo of a forklift is a box.
Fallback: what the robot does when it is unsure
Every physical system needs a defined answer to "what now?" when the model is uncertain, the frame is bad, or the drift alarm is on. Design it explicitly, before the model, and ideally in a layer that does not depend on the neural network at all. A small, boring rule layer looks like this:
from collections import deque
import numpy as np
class FrameMonitor:
"""Tracks simple input statistics and compares them with numbers from the training set."""
def __init__(self, ref, window=300, tol=3.0):
self.ref = ref # {"brightness": (mean, std), ...}
self.win = {k: deque(maxlen=window) for k in ref}
self.tol = tol
def update(self, gray, conf):
self.win["brightness"].append(float(gray.mean()))
self.win["confidence"].append(float(conf))
def drifted(self):
alarms = []
for name, (mu, sd) in self.ref.items():
w = self.win[name]
if len(w) < w.maxlen:
continue # not enough data yet
se = sd / np.sqrt(len(w)) # standard error of the window mean
if abs(np.mean(w) - mu) > self.tol * se:
alarms.append(name)
return alarms
def decide(probs, conf_ok=0.85, conf_floor=0.55):
top = float(probs.max())
if top >= conf_ok:
return "ACT"
if top >= conf_floor:
return "SLOW_AND_RECHECK" # e.g. slow down, take another frame, change viewpoint
return "SAFE_STOP" # stop, hold position, or hand over to a human
The thresholds are placeholders: choose them from your calibrated confidences and the cost of each kind of mistake on a validation set, not by feel. The important part is the structure. There are three outcomes, not two, so the system has a graceful middle ground; the safest action is always available; and the drift monitor can override the model's opinion when the input itself looks wrong. In a warehouse, "uncertain" might mean "flag for a human to look at this pallet later", which is cheap. On a drone near people, "uncertain" must mean "hold position or land". The cost of each decision belongs to the application, not the model.
Pipeline overlap and batching
If you run capture, pre-processing, inference and post-processing one after another, your frame time is the sum. If they overlap, so that the GPU is working on frame n while the CPU prepares frame n+1, it is closer to the slowest stage. Threads work well here because OpenCV, ONNX Runtime and CUDA release Python's global interpreter lock during the heavy calls; for heavier pipelines, look at GStreamer or NVIDIA DeepStream, which are built for this.
For real-time systems, the important rule is that a queue must be short and must drop the oldest frame, not block. A long queue is a hidden latency: the model would be answering about what the camera saw a second ago.
import queue
def capture_loop(camera, frames: queue.Queue, running):
while running.is_set():
frame = camera.read()
if frames.full():
try:
frames.get_nowait() # drop the oldest: fresh beats complete
except queue.Empty:
pass
frames.put(frame)
# frames = queue.Queue(maxsize=2)
# One thread runs capture_loop; another does: get -> preprocess -> infer -> postprocess.
Batching is a different trade. Larger batches raise throughput on a GPU, because the weights are reused across images, but they add latency, since you wait to fill the batch. For a robot reacting to the latest frame, batch size 1 is right. For an offline job, like a camera reviewing a day of recorded footage overnight, a large batch is right. Decide which one you are before you tune.
Updates, rollback and the data flywheel
Once devices are in the field, you need a safe way to change the model. The pattern that works:
- Version everything: the model file, its preprocessing config, its calibration set, its temperature and thresholds, as one bundle with a hash. A model and its preprocessing shipped separately will one day mismatch.
- Sign and verify bundles on the device. A robot that accepts any file is a security hole.
- Roll out gradually. Update a small slice of the fleet first, compare its monitored statistics against the old model on comparable devices, and widen only if they look healthy.
- Keep the previous version on the device and switch back automatically if the new one crashes, fails a self-test on start-up, or trips the drift monitor. Rollback must not need the network.
- Shadow mode is very useful: run the new model beside the old one, log both outputs, act only on the old one, and look at where they disagree.
The data flywheel is what keeps the system improving. The device saves the frames it is least sure about (low confidence, disagreement between the model and a human, or between two models), uploads a small budget of them when a connection is available, people label them, and the next training run includes them. Hard examples are worth far more than random ones. Cap what you store, and record the model version and conditions with each sample.
Privacy of on-device data
Cameras see people, faces, screens, licence plates and private property. The best privacy feature of edge deployment is that raw video need not leave the device at all: process locally and send only results. When you do collect frames for the flywheel, collect as few as you can, blur or drop faces and plates on the device before upload, encrypt at rest and in transit, set retention limits, and tell the people who work near the device what is captured. Check the regulations that apply where the device operates. This is not a place for improvisation, and getting it wrong ends projects.
Debugging: "it works in the notebook, fails on the device"
When a model behaves differently on the device, the cause is rarely the neural network. Work through this checklist in order; it goes from the most likely and cheapest culprit to the least.
- Compare the exact input tensors. Save the tensor fed to the model in the notebook and the one fed on the device for the same image, and diff them. Most problems show up here immediately. Do this before anything else.
- Channel order. PyTorch and PIL use RGB; OpenCV's
cv2.imreadand most cameras via OpenCV give BGR. Swapping channels often reduces accuracy without crashing, which is the worst kind of failure. Convert explicitly withcv2.cvtColor(img, cv2.COLOR_BGR2RGB). - Layout and dtype. Training uses (B, C, H, W) floats; camera pipelines often give (H, W, C) uint8. Check the transpose, the dtype and whether your scaling to [0, 1] happens exactly once (not zero times, not twice).
- Normalisation constants. These are model-specific. Many ImageNet models use mean (0.485, 0.456, 0.406) and std (0.229, 0.224, 0.225); some ViT checkpoints use 0.5 for both; CLIP has its own values. Take them from the checkpoint's own preprocessing config, never from memory.
- Resize method. PIL bilinear, OpenCV
INTER_LINEAR, OpenCVINTER_AREA, and torchvision's resize with or withoutantialiasgive different pixels, especially when shrinking a large frame. Models are sensitive to this. Use the same library and the same interpolation on the device as in training, or build a test that measures the difference on a few hundred images. Also check the geometry: resizing to a square distorts the aspect ratio, while resize-then-centre-crop throws away the edges, and the two must match training. - Unsupported or mismatched ONNX operators. If export or engine build fails, or a layer silently falls back to the CPU, look at the operator list. Update the runtime, change the opset, or rewrite the offending operation using standard ops. Compare intermediate outputs between PyTorch and ONNX Runtime to find the first layer that diverges.
- Frozen shapes. A model exported with a fixed input size or batch of 1 will fail or misbehave on other shapes. Check your
dynamic_axes, and remember positional embeddings are tied to the token count. - Quantisation accuracy drop. Evaluate the FP32 ONNX, the FP16 engine and the INT8 engine on the same held-out set, as separate numbers. If INT8 alone falls, the issue is calibration (wrong data, outliers, sensitive layers). If FP16 falls, look for overflow in activations.
- Device conditions. Compare a saved image from the device's camera with your training images by eye and by statistics: exposure, white balance, compression artefacts, lens blur, frame rate. A camera's auto-exposure can produce something quite unlike your dataset.
- Thermal and power state. If it is fast for a minute and slow afterwards, or accuracy is fine but latency is erratic, check temperatures, clocks and the power mode.
- Runtime versions. Record the versions of the driver, CUDA, TensorRT and ONNX Runtime on the device and in your test environment. Differences here explain a surprising number of "it worked yesterday" reports.
A sturdy way to prevent most of this is a preprocessing parity test that lives in your repository. It takes ten fixed images, runs them through the training preprocessing and the deployment preprocessing, and asserts that the tensors match within a small tolerance. Run it in CI and again on the device, and store a golden set of logits from the reference model: after every export, quantisation or upgrade, run the same images and check that the outputs are close enough.
import numpy as np
def check_parity(train_pre, deploy_pre, images, atol=2e-2):
"""Both callables take one image and return a float32 (3, H, W) array."""
for i, img in enumerate(images):
a, b = train_pre(img), deploy_pre(img)
worst = np.abs(a - b).max()
assert a.shape == b.shape, f"image {i}: shape {a.shape} vs {b.shape}"
assert worst < atol, f"image {i}: max diff {worst:.4f} (channels, resize or norm mismatch?)"
The tolerance needs to allow for the small differences that different resize implementations can legitimately produce. If your differences are larger than that, do not raise the tolerance; find the difference.
Three scenarios, three sets of trade-offs
It helps to see how the pieces combine in different settings. These are imagined scenarios to illustrate the reasoning, not case studies.
The warehouse robot and the damaged box. The camera is fixed to the robot, the lighting is consistent, and boxes look alike. Decisions are not safety-critical: the cost of a false alarm is a human glance. Here a small model at moderate resolution, INT8 on a Jetson-class board, is likely enough, and the interesting design work is in the confidence threshold and the flywheel. Send uncertain "maybe damaged" frames to a person, label them, and retrain monthly. The main drift risk is new packaging, so watch the class distribution and the confidence.
The farm drone counting crops. Images are large, taken from height, and the objects (seedlings) are tiny; the field looks different in every season and in every kind of light. Downsampling the frame would erase the objects, so a naive ViT on the full image hits the quadratic wall. Tile the image into crops and run a small model per tile, or use a windowed or hybrid backbone; a frozen DINOv2-style backbone with a light head is attractive because labelled examples are scarce and change by season. Battery and thermal limits dominate, so quantise, and prefer processing on landing rather than in flight if the mission allows it. Confidence and calibration matter less than counting error, so evaluate on counts.
The factory camera hunting for defects. A fixed camera, controlled light and a very high resolution of small scratches on parts that pass at speed. Defects are rare, so the dataset is heavily imbalanced, and accuracy is a misleading metric: a model that never reports a defect is 99.9% "accurate". Use precision and recall at a chosen threshold, and calibrate that threshold on a validation set that reflects real rates. A missed defect and a false stop have known costs in money, and the threshold should reflect them. Latency is set by the conveyor speed; consider an anomaly-detection approach using frozen features when you have few defect examples. Drift here is things like a new supplier's surface finish or a lamp that dims with age, so log brightness statistics.
What to build next
If you want to make this real, here is a sequence that will teach you more than reading another article.
- Run the ViT demo above, then swap in a pre-trained small ViT or DINOv2 backbone (from
timmortorch.hub), and fine-tune or train a linear head on a tiny dataset of your own, such as photos of two kinds of object on your desk. - Export it to ONNX, verify against PyTorch, and run it in ONNX Runtime. Write the preprocessing parity test.
- Quantise with real calibration images from your own camera. Measure accuracy on a held-out set and report per-class numbers.
- Benchmark with the harness above on every device you can get hold of: a laptop CPU, a phone, and a Jetson-class board if you have one. Include a five-minute soak test.
- Add the input monitor and a three-state decision rule, then deliberately break the input (cover half the lens, turn off the lights, use a blurred video) and check the system reacts the way you designed.
- Write down your rollback plan before you write the update mechanism.
Keep a small log of each step's accuracy, latency and size. The habit of always measuring the three together is what separates a demo from a system.
Where this goes next
We have gone from pixels to tokens to a deployed, monitored model, and in all of it I have leaned on a handful of mathematical ideas without stopping for them: why a softmax behaves as it does, what variance and expectation say about the scale of a dot product, how gradients flow through a residual stack, what quantisation error looks like as a statistic, and why calibration is a statement about probabilities. Those pieces are the same ones that underpin agents, fine-tuning and data pipelines.
In Part 5 of this series we step back and build the mathematical foundations of AI properly: linear algebra, calculus, probability and optimisation, all from first principles and all tied back to code, so that the derivations you skimmed here become things you can do yourself. If you have followed the βd_k argument and the quantisation example in this article, you are already in good shape for it.
Until then, take one model, put it on a device you can hold, and measure it honestly. Everything else in this article is there to help you understand the number you get.
Comments (0)
No comments yet. Be the first to share your thoughts.