Energy-Efficient On-Device ML: Benchmarking Quantization, Pruning, and Runtime Optimizations

Why on-device efficiency matters

Running machine learning models on edge devices impacts battery life, latency, and privacy. This article walks through a practical benchmarking approach to evaluate quantization, pruning, and runtime optimizations for energy-efficient on-device ML. Expect hands-on tips, small code snippets, and concrete steps to reproduce experiments with common tools like PyTorch, TensorFlow Lite, ONNX Runtime, and simple power measurement setups.

Overview of techniques

Quantization – reduce numeric precision (float32 -> int8) to save memory and compute.
Pruning – remove redundant weights or channels to shrink model size and FLOPs.
Runtime optimizations – use vendor runtimes, operator fusion, multithreading, and hardware-specific kernels.

Each technique trades off accuracy, latency, and energy consumption. The goal here is to show how to measure those tradeoffs and how to combine optimizations to get practical wins.

Benchmarking methodology

Pick representative models and datasets: MobileNetV2 or ResNet18 on a small validation split works well for quick experiments.
Define metrics: latency (ms), energy per inference (mJ), memory footprint (MB), and accuracy (top-1, top-5).
Use controlled environment: fixed CPU governor, airplane mode for mobile, and repeated runs to reduce variance.
Use both synthetic and real inputs: synthetic inputs isolate compute; real inputs capture preprocessing cost.

Energy measurement can be done with external power meters, USB power monitors, or built-in counters like RAPL on Intel. For mobile, small external devices such as INA219 or Monsoon power monitor are common in reproducible setups.

Quantization: practical steps

Quantization usually yields large reductions in memory and compute. Try both post-training quantization (PTQ) and quantization-aware training (QAT) when accuracy matters.

Example workflow in PyTorch for PTQ and QAT:

import torch
import torchvision.models as models
from torch.quantization import get_default_qconfig, prepare, convert

# Load model
model = models.mobilenet_v2(pretrained=True).eval()

# PTQ with dynamic quantization for linear layers
model_dyn = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)

# For static PTQ or QAT (simplified)
model_fp32 = models.mobilenet_v2(pretrained=True).eval()
model_fp32.qconfig = get_default_qconfig('fbgemm')
prepare(model_fp32, inplace=True)
# calibrate with a few batches of data
# for inputs, _ in calibration_loader: model_fp32(inputs)
convert(model_fp32, inplace=True)

After quantization, export to ONNX or TFLite for edge runtimes. TFLite supports post-training full integer quantization when given a representative dataset. ONNX Runtime and TensorRT can consume quantized models or apply quantization passes.

Pruning: types and examples

Pruning approaches vary. Structured pruning removes entire filters or channels and often maps better to runtime speedups. Unstructured pruning sparsifies weights but may need sparse kernels to benefit energy and latency.

Example: simple magnitude-based pruning using PyTorch utilities:

import torch
import torch.nn.utils.prune as prune

model = models.resnet18(pretrained=True)
# prune 30 percent of weights in conv layers globally
parameters_to_prune = []
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        parameters_to_prune.append((module, 'weight'))

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.3,
)

# make pruning permanent
for module, _ in parameters_to_prune:
    prune.remove(module, 'weight')

After pruning, fine-tune briefly to recover accuracy. When targeting specific hardware, prefer structured pruning (channel or filter) so runtimes can exploit smaller tensors directly.

Runtime optimizations and runtimes

Choosing the right runtime often matters more than algorithmic tweaks. Experiment with:

TensorFlow Lite for mobile and microcontrollers.
ONNX Runtime with the NNAPI or CoreML delegates on mobile.
TorchScript plus mobile runtime for PyTorch models.
Vendor accelerators like TensorRT, OpenVINO, or NNAPI delegates for hardware-specific kernels.

Also tune thread counts, use operator fusion backends, and enable workspace memory tuning. For many CPU-bound models, setting num threads to the device’s physical cores helps reduce latency and energy per inference.

Measuring energy: practical tips

Example measurement protocol:

Warm up the model and runtime with a set of inferences.
Measure a batch of N inferences and record time and energy.
Compute energy per inference as total energy divided by N, and latency as average per inference.
Repeat with different seeds and inputs to estimate variability.

For code-level measurement on a desktop with RAPL, consider pyRAPL or reading MSRs. For simple USB devices, sample voltage and current at 5 to 100 Hz depending on expected transients.

Combining techniques: a sample pipeline

One effective pipeline to try on a small model:

Start with a compact architecture such as MobileNetV2.
Apply structured pruning to remove channels that contribute little to accuracy.
Fine-tune for a few epochs to restore accuracy.
Perform QAT so the quantized model retains accuracy.
Export to a target runtime (TFLite or ONNX) and benchmark latency and energy.

This order often yields better energy-accuracy tradeoffs than pruning after quantization or random combinations that require additional tuning.

Concrete example: MobileNetV2 to TFLite int8

Below is a summarized flow to convert a trained model to a TFLite int8 file with a representative dataset for calibration.

import tensorflow as tf

# assume saved Keras model
model = tf.keras.applications.MobileNetV2(weights='imagenet')
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

def representative_data_gen():
    for _ in range(100):
        # yield a single sample shaped as [1, 224, 224, 3]
        yield [tf.random.uniform(shape=(1, 224, 224, 3), dtype=tf.float32)]

converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8

tflite_model = converter.convert()
open('mobilenet_v2_int8.tflite', 'wb').write(tflite_model)

After conversion, use the TFLite benchmark tool or a small Python wrapper on-device to measure latency and energy per inference.

Interpreting benchmark results

Key observations to look for:

Memory reduction from quantization often leads to fewer cache misses and lower energy per inference.
Pruning that reduces compute without changing tensor shapes yields limited runtime speedups unless the runtime handles sparse ops.
QAT tends to preserve accuracy better than PTQ on challenging models, at the cost of extra training time.

When comparing setups, visualize tradeoffs with scatter plots: accuracy vs energy and latency vs model size. This makes it clear which optimization provides the most practical benefit for your target device.

Tooling and libraries to try

PyTorch quantization toolkit and NNCF for structured pruning workflows.
TensorFlow Model Optimization Toolkit for pruning and quantization.
ONNX Runtime with quantization tools and hardware delegates.
Vendor SDKs: TensorRT, OpenVINO, Hexagon NN for Qualcomm, and NNAPI/CoreML delegates.

Choosing tools depends on deployment target and whether you can accept model export and conversion complexity for runtime gains.

Practical checklist before deployment

Measure accuracy on a representative validation set after each optimization step.
Record latency and energy under realistic input preprocessing and batching.
Test on the target device with the intended runtime delegate.
Monitor memory usage and thermal throttling on longer runs.
Automate experiments so changes are reproducible and comparable.

Document configurations such as thread count, graph optimization level, and quantization parameters. Small settings can change the energy profile meaningfully.

Closing notes: what to expect

Expect diminishing returns when stacking many small optimizations. Large wins usually come from switching to a more compact architecture, using integer quantization, or leveraging vendor accelerators. Use targeted pruning and QAT when accuracy must be preserved while saving energy.

Try the pipelines described here, measure consistently, and iterate. Practical experimentation on the actual device often reveals surprising interactions between runtime, model structure, and energy use.