2. Model optimization pipeline

Deploying a trained model to an edge device is not a single conversion step; it is a pipeline in which each stage changes the model's size, latency, and accuracy in different ways. A typical edge optimization pipeline consists of four stages: profiling to find where time and memory are actually spent, quantization to reduce numerical precision, pruning to remove redundant parameters, and knowledge distillation to transfer capability from a large teacher model into a smaller student. The order matters, and so does the decision to skip a stage when profiling shows it will not help.

Stage 1: Profiling Before Touching the Model

Profiling answers three questions: which operators dominate latency, where peak memory occurs, and whether the bottleneck is compute-bound or memory-bandwidth-bound. On edge hardware, a model that looks small in parameter count can still be slow because of memory traffic or unsupported operators that fall back to the CPU.

A practical profiling workflow uses a representative input batch and measures per-operator time. For example, with a PyTorch model exported to ONNX and run under ONNX Runtime:

# Enable per-operator profiling in ONNX Runtime
sess_options = ort.SessionOptions()
sess_options.enable_profiling = True
sess = ort.InferenceSession("model.onnx", sess_options,
                            providers=["CPUExecutionProvider"])
sess.run(None, {"input": dummy_input})
prof = sess.end_profiling()   # writes a JSON trace

The resulting trace is parsed to rank operators by total time. A common finding on edge CPUs is that a single Conv or MatMul node consumes 40–60% of latency, while dozens of small element-wise ops consume little time individually but add up through kernel launch overhead. This ranking determines which optimization to apply first: if one heavy operator dominates, quantization of that operator gives the largest return; if many small ops dominate, operator fusion or graph rewriting is more effective.

Stage 2: Quantization

Quantization maps floating-point weights and activations to lower-precision integers, most commonly INT8. The two dominant schemes are post-training quantization (PTQ) and quantization-aware training (QAT).

PTQ requires only a small calibration dataset (typically 100–500 representative samples) and no retraining. It computes scale and zero-point per tensor:

real_value ≈ scale × (quantized_value − zero_point)
scale = (max − min) / (q_max − q_min)

For symmetric weight quantization, zero_point = 0 and the range is [−127, 127]. Activations are usually asymmetric because ReLU outputs are non-negative. Per-channel quantization, where each output channel of a convolution gets its own scale, generally preserves accuracy better than per-tensor quantization and is supported by most edge runtimes.

QAT inserts fake-quantization nodes into the training graph so the model learns to tolerate rounding error. It typically recovers 1–3% accuracy compared to PTQ on sensitive models such as detection networks, at the cost of a full fine-tuning cycle.

After quantization, verify that the target runtime actually supports the quantized operators. An INT8 model that falls back to FP32 for unsupported ops can be slower than the original.

Stage 3: Pruning

Pruning removes parameters that contribute little to the output. Unstructured pruning zeroes individual weights and produces sparse tensors, which only help if the runtime and hardware support sparse computation. Structured pruning removes entire filters, channels, or attention heads, producing a dense smaller model that runs faster on any hardware.

A standard structured pruning loop for a convolutional network:

1. Train the model to convergence.
2. Compute an importance score per filter, e.g. the L1 norm of its weights.
3. Remove the lowest-scoring fraction (e.g. 20%) across all layers.
4. Fine-tune for a few epochs to recover accuracy.
5. Repeat until the target latency or parameter budget is reached.

Pruning ratios should be set per layer, not globally. Early layers and the first/last layers are usually more sensitive; aggressive pruning there damages accuracy disproportionately. A common heuristic is to prune deeper layers more heavily and keep the first convolution and the classifier head nearly intact.

Stage 4: Knowledge Distillation

Knowledge distillation trains a small student model to mimic a large teacher model. Instead of learning only from hard labels, the student learns from the teacher's soft probability distribution, which carries information about class similarity.

The standard loss combines a distillation term and a task term:

L = α · CE(student_logits, hard_labels)
  + (1 − α) · T² · KL(softmax(student_logits / T),
                        softmax(teacher_logits / T))

Here T is the temperature (commonly 2–5) and α balances the two terms. The T² factor keeps gradient magnitudes comparable when temperature changes. Distillation is especially useful when the target device cannot run the teacher but the teacher's accuracy is needed as a ceiling.

Combining the Stages

The four stages are not independent. A practical order for a CNN classifier on an edge NPU is: profile → structured prune → fine-tune → quantize (PTQ, fall back to QAT if accuracy drops too far) → re-profile to confirm the latency gain. Distillation can be applied during the fine-tuning step, using the original unpruned model as teacher.

The numbers in the diagram are illustrative for a mid-sized CNN and will vary by architecture, hardware, and calibration set. The key engineering discipline is to measure after every stage rather than assuming the pipeline improves everything. A stage that does not reduce the profiled bottleneck should be reverted, not kept for its own sake.

Verification Checklist

• Re-profile after each stage and confirm the dominant operator actually became faster.
• Validate accuracy on a held-out set, not the calibration set used for quantization.
• Check that all quantized operators are supported by the target runtime; watch for CPU fallback.
• Compare end-to-end latency at the production batch size, not a benchmark batch size.
• Record the configuration of each stage (pruning ratio, calibration size, temperature) so results are reproducible.