On-device inference runs ML models directly on phones, IoT devices, and embedded systems instead of the cloud. This enables real-time responses, works offline, and keeps data private. But edge devices have limited compute, memory, and battery - so models must be optimized carefully.

Where It Is Used

On-device ML powers everyday experiences:

Mobile apps: Camera filters, keyboard prediction, voice assistants
Wearables: Heart rate analysis, activity recognition
Smart home: Wake word detection, local automation
Automotive: Driver monitoring, ADAS features
Industrial IoT: Predictive maintenance, quality inspection
Retail: Barcode scanning, inventory tracking
Healthcare: Real-time monitoring, diagnostic aids

Device Constraints

Constraint	Typical Limits	Impact
Memory	50MB - 2GB RAM	Model size limits
Compute	0.5 - 8 GFLOPS	Inference latency
Storage	10MB - 100MB for model	Compression needed
Battery	5-10mAh per inference	Usage limits
Thermal	Sustained vs burst	Throttling
Network	Often unavailable	Must work offline

What You Do in Practice

Profile target device - CPU, GPU, NPU capabilities
Set latency budget - User experience requirements
Compress model - Quantization, pruning, distillation
Choose runtime - TFLite, ONNX, CoreML, etc.
Optimize for hardware - Use accelerators
Benchmark extensively - Real device testing
Handle edge cases - Low battery, background mode
Update strategy - OTA model updates

Model Optimization Pipeline

import torch
import torch.nn as nn
import numpy as np
from typing import Tuple, Dict, Any

class EdgeModelOptimizer:
    """Comprehensive optimization for edge deployment."""

    def __init__(self, model: nn.Module, target_latency_ms: float = 50):
        self.model = model
        self.target_latency_ms = target_latency_ms
        self.optimization_history = []

    def full_optimization_pipeline(self, 
                                   calibration_data: torch.Tensor,
                                   val_data: Tuple[torch.Tensor, torch.Tensor]
                                   ) -> nn.Module:
        """Run complete optimization pipeline."""
        original_size = self.get_model_size()
        original_latency = self.measure_latency()
        original_accuracy = self.evaluate(val_data)

        print(f"Original: {original_size:.2f}MB, {original_latency:.1f}ms, {original_accuracy:.3f}")

        # Step 1: Prune unimportant weights
        model = self.structured_pruning(sparsity=0.3)

        # Step 2: Knowledge distillation (if teacher available)
        # model = self.distill(teacher_model, train_data)

        # Step 3: Quantization
        model = self.quantize_dynamic()

        # Step 4: Operator fusion
        model = self.fuse_operations()

        final_size = self.get_model_size()
        final_latency = self.measure_latency()
        final_accuracy = self.evaluate(val_data)

        print(f"Optimized: {final_size:.2f}MB, {final_latency:.1f}ms, {final_accuracy:.3f}")
        print(f"Reduction: {(1-final_size/original_size)*100:.1f}% size, "
              f"{(1-final_latency/original_latency)*100:.1f}% latency")

        return model

    def quantize_dynamic(self) -> nn.Module:
        """Apply dynamic quantization."""
        quantized = torch.quantization.quantize_dynamic(
            self.model,
            {nn.Linear, nn.LSTM, nn.GRU},
            dtype=torch.qint8
        )
        self.model = quantized
        return quantized

    def quantize_static(self, calibration_data: torch.Tensor) -> nn.Module:
        """Apply static quantization with calibration."""
        self.model.eval()

        # Fuse modules
        self.model = torch.quantization.fuse_modules(
            self.model, [['conv', 'bn', 'relu']]
        )

        # Prepare for quantization
        self.model.qconfig = torch.quantization.get_default_qconfig('qnnpack')
        torch.quantization.prepare(self.model, inplace=True)

        # Calibrate with representative data
        with torch.no_grad():
            for batch in calibration_data:
                self.model(batch)

        # Convert to quantized
        torch.quantization.convert(self.model, inplace=True)

        return self.model

    def structured_pruning(self, sparsity: float = 0.3) -> nn.Module:
        """Remove structured components (channels/filters)."""
        import torch.nn.utils.prune as prune

        for name, module in self.model.named_modules():
            if isinstance(module, nn.Conv2d):
                prune.ln_structured(
                    module, name='weight', 
                    amount=sparsity, n=2, dim=0
                )
                prune.remove(module, 'weight')
            elif isinstance(module, nn.Linear):
                prune.l1_unstructured(module, name='weight', amount=sparsity)
                prune.remove(module, 'weight')

        return self.model

    def fuse_operations(self) -> nn.Module:
        """Fuse operations for faster inference."""
        # Conv + BN + ReLU fusion
        for name, module in self.model.named_children():
            if hasattr(module, 'fuse'):
                module.fuse()

        return self.model

    def get_model_size(self) -> float:
        """Get model size in MB."""
        param_size = sum(
            p.nelement() * p.element_size() 
            for p in self.model.parameters()
        )
        buffer_size = sum(
            b.nelement() * b.element_size() 
            for b in self.model.buffers()
        )
        return (param_size + buffer_size) / (1024 ** 2)

    def measure_latency(self, input_shape: Tuple = (1, 3, 224, 224),
                        warmup: int = 10, runs: int = 100) -> float:
        """Measure inference latency in ms."""
        self.model.eval()
        dummy_input = torch.randn(*input_shape)

        # Warmup
        with torch.no_grad():
            for _ in range(warmup):
                self.model(dummy_input)

        # Measure
        import time
        times = []
        with torch.no_grad():
            for _ in range(runs):
                start = time.perf_counter()
                self.model(dummy_input)
                times.append((time.perf_counter() - start) * 1000)

        return np.median(times)

    def evaluate(self, val_data: Tuple[torch.Tensor, torch.Tensor]) -> float:
        """Evaluate model accuracy."""
        self.model.eval()
        X, y = val_data
        with torch.no_grad():
            outputs = self.model(X)
            preds = outputs.argmax(dim=1)
            accuracy = (preds == y).float().mean().item()
        return accuracy

Quantization Strategies

import torch
import torch.nn as nn
from torch.quantization import QuantStub, DeQuantStub

class QuantizableModel(nn.Module):
    """Model prepared for quantization."""

    def __init__(self, original_model: nn.Module):
        super().__init__()
        self.quant = QuantStub()
        self.model = original_model
        self.dequant = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.model(x)
        x = self.dequant(x)
        return x

class QuantizationAwareTraining:
    """QAT for better accuracy with quantization."""

    def __init__(self, model: nn.Module, backend: str = 'qnnpack'):
        self.model = model
        self.backend = backend

    def prepare_qat(self):
        """Prepare model for quantization-aware training."""
        self.model.train()

        # Set quantization config
        self.model.qconfig = torch.quantization.get_default_qat_qconfig(
            self.backend
        )

        # Fuse operations
        self.model = torch.quantization.fuse_modules(
            self.model, 
            self._get_fusion_patterns()
        )

        # Prepare for QAT
        torch.quantization.prepare_qat(self.model, inplace=True)

        return self.model

    def _get_fusion_patterns(self):
        """Identify fusable patterns in model."""
        patterns = []
        # Find Conv-BN-ReLU patterns
        modules = list(self.model.named_modules())
        for i, (name, module) in enumerate(modules[:-2]):
            if isinstance(module, nn.Conv2d):
                next_name, next_module = modules[i+1]
                third_name, third_module = modules[i+2]
                if isinstance(next_module, nn.BatchNorm2d) and \
                   isinstance(third_module, nn.ReLU):
                    patterns.append([name, next_name, third_name])
        return patterns

    def train_qat(self, train_loader, epochs: int = 5):
        """Train with quantization awareness."""
        optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)
        criterion = nn.CrossEntropyLoss()

        for epoch in range(epochs):
            self.model.train()

            for batch_x, batch_y in train_loader:
                optimizer.zero_grad()
                outputs = self.model(batch_x)
                loss = criterion(outputs, batch_y)
                loss.backward()
                optimizer.step()

        return self.model

    def convert_to_quantized(self):
        """Convert to fully quantized model."""
        self.model.eval()
        quantized = torch.quantization.convert(self.model, inplace=False)
        return quantized

class MixedPrecisionInference:
    """Different precision for different layers."""

    def __init__(self, model: nn.Module):
        self.model = model
        self.layer_precision = {}

    def analyze_sensitivity(self, val_data, test_precisions=[8, 16, 32]):
        """Find optimal precision per layer."""
        baseline_accuracy = self._evaluate(val_data)

        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                sensitivities = {}

                for bits in test_precisions:
                    # Quantize just this layer
                    self._quantize_layer(name, bits)
                    accuracy = self._evaluate(val_data)
                    accuracy_drop = baseline_accuracy - accuracy
                    sensitivities[bits] = accuracy_drop
                    self._restore_layer(name)

                # Choose minimum precision that maintains accuracy
                for bits in test_precisions:
                    if sensitivities[bits] < 0.01:  # < 1% drop
                        self.layer_precision[name] = bits
                        break
                else:
                    self.layer_precision[name] = 32

        return self.layer_precision

    def _evaluate(self, val_data):
        """Evaluate model accuracy."""
        self.model.eval()
        X, y = val_data
        with torch.no_grad():
            outputs = self.model(X)
            preds = outputs.argmax(dim=1)
        return (preds == y).float().mean().item()

Runtime Selection

from enum import Enum
from typing import Optional
import platform

class MLRuntime(Enum):
    TFLITE = "tflite"
    ONNX = "onnx"
    COREML = "coreml"
    PYTORCH_MOBILE = "pytorch_mobile"
    TENSORRT = "tensorrt"

class RuntimeSelector:
    """Select optimal runtime for target device."""

    RUNTIME_CAPABILITIES = {
        MLRuntime.TFLITE: {
            'platforms': ['android', 'ios', 'linux', 'windows'],
            'accelerators': ['cpu', 'gpu', 'npu', 'dsp'],
            'quantization': ['int8', 'float16'],
            'model_size_limit_mb': 100
        },
        MLRuntime.COREML: {
            'platforms': ['ios', 'macos'],
            'accelerators': ['cpu', 'gpu', 'neural_engine'],
            'quantization': ['int8', 'float16'],
            'model_size_limit_mb': 200
        },
        MLRuntime.ONNX: {
            'platforms': ['android', 'ios', 'linux', 'windows'],
            'accelerators': ['cpu', 'gpu'],
            'quantization': ['int8', 'float16', 'float32'],
            'model_size_limit_mb': 500
        },
        MLRuntime.TENSORRT: {
            'platforms': ['linux', 'jetson'],
            'accelerators': ['gpu'],
            'quantization': ['int8', 'float16', 'float32'],
            'model_size_limit_mb': 1000
        }
    }

    def select_runtime(self, target_platform: str,
                       model_size_mb: float,
                       required_accelerator: str = 'cpu',
                       quantization: str = 'float32') -> MLRuntime:
        """Select best runtime for requirements."""
        candidates = []

        for runtime, caps in self.RUNTIME_CAPABILITIES.items():
            if target_platform in caps['platforms'] and \
               required_accelerator in caps['accelerators'] and \
               quantization in caps['quantization'] and \
               model_size_mb <= caps['model_size_limit_mb']:
                candidates.append(runtime)

        if not candidates:
            raise ValueError(f"No runtime supports requirements")

        # Priority: native platform runtime > general purpose
        if target_platform == 'ios' and MLRuntime.COREML in candidates:
            return MLRuntime.COREML
        if target_platform == 'android' and MLRuntime.TFLITE in candidates:
            return MLRuntime.TFLITE
        if required_accelerator == 'gpu' and MLRuntime.TENSORRT in candidates:
            return MLRuntime.TENSORRT

        return candidates[0]

class TFLiteConverter:
    """Convert PyTorch model to TensorFlow Lite."""

    def convert(self, pytorch_model: nn.Module,
                input_shape: Tuple,
                quantize: bool = True) -> bytes:
        """Convert and optionally quantize model."""
        import torch.onnx
        import onnx
        from onnx_tf.backend import prepare
        import tensorflow as tf

        # PyTorch -> ONNX
        dummy_input = torch.randn(*input_shape)
        onnx_path = "temp_model.onnx"

        torch.onnx.export(
            pytorch_model,
            dummy_input,
            onnx_path,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={'input': {0: 'batch_size'}}
        )

        # ONNX -> TensorFlow
        onnx_model = onnx.load(onnx_path)
        tf_rep = prepare(onnx_model)
        tf_path = "temp_model_tf"
        tf_rep.export_graph(tf_path)

        # TensorFlow -> TFLite
        converter = tf.lite.TFLiteConverter.from_saved_model(tf_path)

        if quantize:
            converter.optimizations = [tf.lite.Optimize.DEFAULT]
            converter.target_spec.supported_types = [tf.float16]

        tflite_model = converter.convert()

        return tflite_model

Battery-Aware Inference

import time
from dataclasses import dataclass
from typing import Callable

@dataclass
class InferenceConfig:
    """Configuration for adaptive inference."""
    max_latency_ms: float = 100
    battery_threshold: float = 0.2  # 20%
    thermal_threshold_c: float = 40
    use_accelerator: bool = True

class BatteryAwareInference:
    """Adapt inference based on device state."""

    def __init__(self, models: Dict[str, nn.Module], config: InferenceConfig):
        """
        Initialize with multiple model variants.

        Args:
            models: Dict with 'full', 'medium', 'lite' variants
            config: Inference configuration
        """
        self.models = models
        self.config = config
        self.current_model = 'full'

    def get_device_state(self) -> Dict:
        """Get current device state (mock implementation)."""
        # In production, use platform-specific APIs
        return {
            'battery_level': 0.75,
            'is_charging': False,
            'temperature_c': 35,
            'memory_available_mb': 512
        }

    def select_model(self) -> nn.Module:
        """Select model variant based on device state."""
        state = self.get_device_state()

        # Low battery - use lite model
        if state['battery_level'] < self.config.battery_threshold:
            self.current_model = 'lite'

        # Thermal throttling - reduce to medium
        elif state['temperature_c'] > self.config.thermal_threshold_c:
            self.current_model = 'medium'

        # Good conditions - use full model
        else:
            self.current_model = 'full'

        return self.models[self.current_model]

    def infer(self, input_tensor: torch.Tensor) -> torch.Tensor:
        """Run inference with adaptive model selection."""
        model = self.select_model()
        model.eval()

        with torch.no_grad():
            output = model(input_tensor)

        return output

class BatchedInference:
    """Batch requests for efficient inference."""

    def __init__(self, model: nn.Module, 
                 max_batch_size: int = 8,
                 max_wait_ms: float = 50):
        self.model = model
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.pending_requests = []

    def add_request(self, input_tensor: torch.Tensor, 
                    callback: Callable):
        """Add inference request to batch."""
        self.pending_requests.append({
            'input': input_tensor,
            'callback': callback,
            'timestamp': time.time()
        })

        # Process if batch full or oldest request waited too long
        if self._should_process():
            self._process_batch()

    def _should_process(self) -> bool:
        """Check if batch should be processed."""
        if len(self.pending_requests) >= self.max_batch_size:
            return True

        if self.pending_requests:
            oldest = self.pending_requests[0]['timestamp']
            if (time.time() - oldest) * 1000 > self.max_wait_ms:
                return True

        return False

    def _process_batch(self):
        """Process all pending requests."""
        if not self.pending_requests:
            return

        # Stack inputs
        inputs = torch.stack([r['input'] for r in self.pending_requests])

        # Run batch inference
        self.model.eval()
        with torch.no_grad():
            outputs = self.model(inputs)

        # Call callbacks with results
        for i, request in enumerate(self.pending_requests):
            request'callback'

        self.pending_requests.clear()

Evaluation Metrics

Metric	Target	How to Measure
Latency	< 50ms for real-time	Benchmark on device
Model size	< 50MB typical	torch.save size
Memory peak	< RAM budget	Profiler
Power	mJ per inference	Power monitor
Accuracy drop	< 1-2% from full	Val set comparison
Startup time	< 1s	First inference

Failure Modes and Debugging Checklist

[ ] Accuracy degradation: Too aggressive optimization
[ ] Latency spikes: Thermal throttling, GC
[ ] Memory overflow: Model too large
[ ] Battery drain: Continuous inference
[ ] Incompatible ops: Unsupported operators
[ ] Device fragmentation: Different hardware
[ ] Update failures: OTA model updates

Production Checklist

Before deployment: - [ ] Benchmark on target devices - [ ] Test battery impact - [ ] Validate accuracy after optimization - [ ] Test offline operation - [ ] Define fallback behaviour

After deployment: - [ ] Monitor latency distribution - [ ] Track battery reports - [ ] A/B test model updates - [ ] Device compatibility testing - [ ] User experience metrics

Tools and Stack

Category	Tools
Conversion	TFLite, ONNX, CoreML Tools
Optimization	TensorRT, NNAPI, Metal
Profiling	Android Profiler, Instruments
Deployment	Firebase ML, AWS IoT Greengrass
Frameworks	MediaPipe, ML Kit

Example Mini-Case

Scenario: Real-time object detection on mobile.

Approach: - MobileNetV3 + SSD architecture - INT8 quantization with QAT - TFLite with GPU delegate - Adaptive resolution based on battery

Results: - 15ms latency on mid-range phone - 4.2MB model size - < 1% accuracy drop from FP32 - 40% power reduction vs CPU

Index: ML Field Guide
Model Compression
Privacy-Preserving ML
Object Detection

Try it live

Everything above runs in your browser — open Personal AI Assistants: On-Device vs Cloud LLM Inference Simulator and change the parameters while it is running. Nothing is installed, nothing is uploaded, the whole model lives in one tab.

▶ Open Personal AI Assistants: On-Device vs Cloud LLM Inference Simulator simulation

On-Device Inference