Edge Computing for AI: Running Models Where They Matter

import onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic, QuantType
 
# Dynamic quantization — the easiest starting point
# Converts weights to INT8, keeps activations in FP32
# Typically 2-4x speedup with <2% accuracy loss
quantize_dynamic(
    model_input="model_fp32.onnx",
    model_output="model_int8.onnx",
    weight_type=QuantType.QInt8
)
 
# Static quantization — better performance, more work
# Requires a calibration dataset to determine optimal scale factors
from onnxruntime.quantization import quantize_static, CalibrationDataReader
 
class MyCalibrationReader(CalibrationDataReader):
    """Feed representative data to determine quantization parameters."""
 
    def __init__(self, calibration_data: list):
        self.data = iter(calibration_data)
 
    def get_next(self):
        try:
            sample = next(self.data)
            return {"input": sample}
        except StopIteration:
            return None
 
# Use 100-500 representative samples
calibration_reader = MyCalibrationReader(calibration_samples)
 
quantize_static(
    model_input="model_fp32.onnx",
    model_output="model_int8_static.onnx",
    calibration_data_reader=calibration_reader,
)

import torch
import torch.nn.utils.prune as prune
 
# Unstructured pruning — remove individual weights
# Good for: general size reduction
model = load_your_model()
 
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        # Remove 30% of weights with smallest magnitude
        prune.l1_unstructured(module, name='weight', amount=0.3)
 
# Structured pruning — remove entire neurons/channels
# Good for: actual speedup (not just smaller files)
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        # Remove 20% of output channels
        prune.ln_structured(
            module, name='weight', amount=0.2, n=2, dim=0
        )
 
# Make pruning permanent (remove the masks)
for name, module in model.named_modules():
    if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
        prune.remove(module, 'weight')
 
# Fine-tune for a few epochs to recover accuracy
# This step matters — don't skip it
fine_tune(model, train_loader, epochs=5, lr=1e-5)

import torch
import torch.nn.functional as F
 
def distillation_loss(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    labels: torch.Tensor,
    temperature: float = 3.0,
    alpha: float = 0.7
) -> torch.Tensor:
    """
    Combined loss: soft targets from teacher + hard targets from labels.
 
    temperature: Higher = softer probability distribution = more knowledge
                 transfer. 3-5 works well for most cases.
    alpha: Balance between teacher knowledge and ground truth.
           0.7 means 70% teacher, 30% hard labels.
    """
    # Soft loss — learn from teacher's probability distribution
    soft_loss = F.kl_div(
        F.log_softmax(student_logits / temperature, dim=-1),
        F.softmax(teacher_logits / temperature, dim=-1),
        reduction='batchmean'
    ) * (temperature ** 2)
 
    # Hard loss — standard cross-entropy with ground truth
    hard_loss = F.cross_entropy(student_logits, labels)
 
    return alpha * soft_loss + (1 - alpha) * hard_loss
 
# Training loop
teacher_model.eval()  # Teacher is frozen
student_model.train()
 
for batch in train_loader:
    inputs, labels = batch
 
    with torch.no_grad():
        teacher_logits = teacher_model(inputs)
 
    student_logits = student_model(inputs)
 
    loss = distillation_loss(student_logits, teacher_logits, labels)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

import torch
import onnx
 
# Export PyTorch model to ONNX
model = load_your_trained_model()
model.eval()
 
# Create dummy input matching your model's expected shape
dummy_input = torch.randn(1, 3, 224, 224)  # batch, channels, H, W
 
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=True,
    opset_version=17,             # Use latest stable opset
    do_constant_folding=True,     # Optimize constant expressions
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={                 # Allow variable batch size
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
 
# Verify the exported model
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print(f"Model exported: {onnx_model.graph.input[0].type}")

// TypeScript — Node.js edge server or Electron app
import * as ort from 'onnxruntime-node';
 
interface PredictionResult {
  label: string;
  confidence: number;
  inferenceTimeMs: number;
}
 
async function runEdgeInference(
  modelPath: string,
  inputData: Float32Array
): Promise<PredictionResult> {
  // Create session with optimization flags
  const session = await ort.InferenceSession.create(modelPath, {
    executionProviders: ['CUDAExecutionProvider', 'CPUExecutionProvider'],
    graphOptimizationLevel: 'all',
    enableCpuMemArena: true,
    enableMemPattern: true,
  });
 
  // Create input tensor
  const inputTensor = new ort.Tensor('float32', inputData, [1, 3, 224, 224]);
 
  // Run inference and measure time
  const start = performance.now();
  const results = await session.run({ input: inputTensor });
  const inferenceTime = performance.now() - start;
 
  // Process output
  const output = results.output.data as Float32Array;
  const maxIdx = output.indexOf(Math.max(...output));
 
  return {
    label: LABELS[maxIdx],
    confidence: output[maxIdx],
    inferenceTimeMs: Math.round(inferenceTime * 100) / 100,
  };
}
 
// Session caching — don't reload the model on every request!
const sessionCache = new Map<string, ort.InferenceSession>();
 
async function getOrCreateSession(
  modelPath: string
): Promise<ort.InferenceSession> {
  if (!sessionCache.has(modelPath)) {
    const session = await ort.InferenceSession.create(modelPath, {
      executionProviders: ['CPUExecutionProvider'],
      graphOptimizationLevel: 'all',
      intraOpNumThreads: 4,       // Tune for your edge device's cores
      interOpNumThreads: 2,
    });
    sessionCache.set(modelPath, session);
  }
  return sessionCache.get(modelPath)!;
}

import tensorrt as trt
import numpy as np
 
# Convert ONNX to TensorRT engine
def build_tensorrt_engine(
    onnx_path: str,
    engine_path: str,
    precision: str = "fp16",      # "fp32", "fp16", or "int8"
    max_batch_size: int = 1,
    workspace_gb: float = 1.0
) -> None:
    """Build an optimized TensorRT engine from an ONNX model."""
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    parser = trt.OnnxParser(network, logger)
 
    # Parse ONNX model
    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(f"Parse error: {parser.get_error(i)}")
            raise RuntimeError("Failed to parse ONNX model")
 
    # Configure builder
    config = builder.create_builder_config()
    config.set_memory_pool_limit(
        trt.MemoryPoolType.WORKSPACE,
        int(workspace_gb * (1 << 30))
    )
 
    # Set precision
    if precision == "fp16":
        if builder.platform_has_fast_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
            print("Using FP16 precision")
        else:
            print("FP16 not supported, falling back to FP32")
    elif precision == "int8":
        if builder.platform_has_fast_int8:
            config.set_flag(trt.BuilderFlag.INT8)
            # INT8 requires calibration — see quantization section
            config.int8_calibrator = create_calibrator(calibration_data)
            print("Using INT8 precision")
 
    # Set dynamic shapes
    profile = builder.create_optimization_profile()
    input_shape = network.get_input(0).shape
    profile.set_shape(
        network.get_input(0).name,
        min=(1, *input_shape[1:]),
        opt=(max_batch_size // 2, *input_shape[1:]),
        max=(max_batch_size, *input_shape[1:])
    )
    config.add_optimization_profile(profile)
 
    # Build engine (this takes minutes — do it once, save the result)
    print("Building TensorRT engine (this takes a while)...")
    serialized_engine = builder.build_serialized_network(network, config)
 
    with open(engine_path, 'wb') as f:
        f.write(serialized_engine)
 
    print(f"Engine saved to {engine_path}")

#!/usr/bin/env python3
"""Edge inference service for NVIDIA Jetson devices."""
 
import asyncio
import time
import logging
from pathlib import Path
from dataclasses import dataclass
from contextlib import asynccontextmanager
 
import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
 
logger = logging.getLogger(__name__)
 
@dataclass
class ModelConfig:
    model_path: str
    input_name: str = "input"
    input_shape: tuple = (1, 3, 224, 224)
    device: str = "cuda"     # "cuda" for Jetson GPU, "cpu" for fallback
    num_threads: int = 4
 
class InferenceRequest(BaseModel):
    data: list[float]
 
class InferenceResponse(BaseModel):
    predictions: list[float]
    inference_ms: float
    device: str
    model_version: str
 
class EdgeInferenceEngine:
    """Manages model lifecycle and inference on edge devices."""
 
    def __init__(self, config: ModelConfig):
        self.config = config
        self.session = None
        self.model_version = "unknown"
        self._inference_count = 0
        self._total_inference_ms = 0.0
 
    def load(self) -> None:
        """Load model with appropriate execution provider."""
        providers = []
 
        if self.config.device == "cuda":
            providers.append(('CUDAExecutionProvider', {
                'device_id': 0,
                'arena_extend_strategy': 'kSameAsRequested',
                'gpu_mem_limit': 512 * 1024 * 1024,  # 512MB — leave room
                'cudnn_conv_algo_search': 'HEURISTIC',
            }))
 
        providers.append(('CPUExecutionProvider', {
            'arena_extend_strategy': 'kSameAsRequested',
        }))
 
        session_options = ort.SessionOptions()
        session_options.graph_optimization_level = (
            ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        )
        session_options.intra_op_num_threads = self.config.num_threads
        session_options.enable_mem_pattern = True
 
        self.session = ort.InferenceSession(
            self.config.model_path,
            sess_options=session_options,
            providers=providers
        )
 
        actual_provider = self.session.get_providers()[0]
        logger.info(f"Model loaded on {actual_provider}")
 
        # Warm-up run — first inference is always slow
        dummy = np.random.randn(*self.config.input_shape).astype(np.float32)
        self.session.run(None, {self.config.input_name: dummy})
        logger.info("Warm-up inference complete")
 
    def predict(self, input_data: np.ndarray) -> tuple[np.ndarray, float]:
        """Run inference and return (output, time_ms)."""
        if self.session is None:
            raise RuntimeError("Model not loaded")
 
        start = time.perf_counter()
        outputs = self.session.run(
            None,
            {self.config.input_name: input_data}
        )
        elapsed_ms = (time.perf_counter() - start) * 1000
 
        self._inference_count += 1
        self._total_inference_ms += elapsed_ms
 
        return outputs[0], elapsed_ms
 
    @property
    def avg_inference_ms(self) -> float:
        if self._inference_count == 0:
            return 0.0
        return self._total_inference_ms / self._inference_count
 
# --- FastAPI service ---
engine = EdgeInferenceEngine(ModelConfig(
    model_path="/models/anomaly_detector_int8.onnx",
    device="cuda"
))
 
@asynccontextmanager
async def lifespan(app: FastAPI):
    engine.load()
    yield
    logger.info(f"Shutting down. Avg inference: {engine.avg_inference_ms:.2f}ms")
 
app = FastAPI(title="Edge Inference Service", lifespan=lifespan)
 
@app.post("/predict", response_model=InferenceResponse)
async def predict(request: InferenceRequest):
    try:
        input_array = np.array(
            request.data, dtype=np.float32
        ).reshape(engine.config.input_shape)
 
        predictions, inference_ms = engine.predict(input_array)
 
        return InferenceResponse(
            predictions=predictions.flatten().tolist(),
            inference_ms=round(inference_ms, 2),
            device=engine.session.get_providers()[0],
            model_version=engine.model_version,
        )
    except Exception as e:
        logger.error(f"Inference failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))
 
@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "inference_count": engine._inference_count,
        "avg_inference_ms": round(engine.avg_inference_ms, 2),
        "provider": engine.session.get_providers()[0] if engine.session else None,
    }

# Raspberry Pi deployment — CPU-only or with Coral USB Accelerator
# Key constraints:
#   - 4-8GB RAM (model must fit with room for the OS)
#   - ARM Cortex CPU (no GPU)
#   - Thermal throttling above ~80°C
 
import onnxruntime as ort
import numpy as np
 
def create_pi_session(model_path: str) -> ort.InferenceSession:
    """Create an optimized session for Raspberry Pi."""
    options = ort.SessionOptions()
 
    # Critical: limit threads to physical cores
    # Pi 5 has 4 cores — using more causes thrashing
    options.intra_op_num_threads = 4
    options.inter_op_num_threads = 1
 
    # Enable all graph optimizations
    options.graph_optimization_level = (
        ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    )
 
    # Memory optimization — important on 4GB Pi
    options.enable_cpu_mem_arena = True
    options.enable_mem_pattern = True
 
    session = ort.InferenceSession(
        model_path,
        sess_options=options,
        providers=['CPUExecutionProvider']
    )
 
    return session
 
# Pro tip: monitor temperature and back off if throttling
def get_cpu_temp() -> float:
    """Read Pi CPU temperature."""
    with open('/sys/class/thermal/thermal_zone0/temp', 'r') as f:
        return float(f.read().strip()) / 1000.0
 
def should_throttle() -> bool:
    """Check if we should reduce inference frequency."""
    temp = get_cpu_temp()
    if temp > 75.0:
        logger.warning(f"CPU temp {temp}°C — throttling inference")
        return True
    return False

// Browser inference with ONNX Runtime Web
// Works in all modern browsers via WebAssembly
 
import * as ort from 'onnxruntime-web';
 
// Configure WASM backend
ort.env.wasm.numThreads = navigator.hardwareConcurrency || 4;
ort.env.wasm.simd = true;
 
interface BrowserInferenceConfig {
  modelUrl: string;            // URL or path to .onnx file
  executionProvider: 'wasm' | 'webgl' | 'webgpu';
  inputShape: number[];
}
 
class BrowserInferenceEngine {
  private session: ort.InferenceSession | null = null;
  private config: BrowserInferenceConfig;
 
  constructor(config: BrowserInferenceConfig) {
    this.config = config;
  }
 
  async initialize(): Promise<void> {
    // Load model — this downloads and compiles the WASM module
    // Show a loading indicator because this takes 1-5 seconds
    this.session = await ort.InferenceSession.create(
      this.config.modelUrl,
      {
        executionProviders: [this.config.executionProvider],
        graphOptimizationLevel: 'all',
      }
    );
 
    // Warm-up run
    const dummy = new Float32Array(
      this.config.inputShape.reduce((a, b) => a * b, 1)
    );
    const tensor = new ort.Tensor('float32', dummy, this.config.inputShape);
    await this.session.run({ input: tensor });
 
    console.log('Model loaded and warmed up');
  }
 
  async predict(inputData: Float32Array): Promise<{
    result: Float32Array;
    timeMs: number;
  }> {
    if (!this.session) throw new Error('Model not initialized');
 
    const tensor = new ort.Tensor(
      'float32', inputData, this.config.inputShape
    );
 
    const start = performance.now();
    const output = await this.session.run({ input: tensor });
    const timeMs = performance.now() - start;
 
    const resultKey = this.session.outputNames[0];
    return {
      result: output[resultKey].data as Float32Array,
      timeMs: Math.round(timeMs * 100) / 100,
    };
  }
 
  dispose(): void {
    this.session?.release();
    this.session = null;
  }
}
 
// Usage in a React component
// const engine = new BrowserInferenceEngine({
//   modelUrl: '/models/classifier_int8.onnx',
//   executionProvider: 'wasm',
//   inputShape: [1, 3, 224, 224],
// });
// await engine.initialize();
// const { result, timeMs } = await engine.predict(preprocessedImage);

import time
import json
import logging
from dataclasses import dataclass, field, asdict
from collections import deque
from typing import Optional
 
@dataclass
class InferenceMetric:
    timestamp: float
    inference_ms: float
    input_hash: str              # For detecting data drift
    prediction: list[float]
    confidence: float
    model_version: str
    device_id: str
    cpu_temp: Optional[float] = None
    memory_usage_mb: Optional[float] = None
 
class EdgeTelemetryCollector:
    """Collects and batches edge inference metrics for upload."""
 
    def __init__(
        self,
        device_id: str,
        buffer_size: int = 1000,
        upload_interval_s: int = 300,  # Upload every 5 minutes
    ):
        self.device_id = device_id
        self.buffer: deque[InferenceMetric] = deque(maxlen=buffer_size)
        self.upload_interval = upload_interval_s
        self._last_upload = time.time()
 
        # Rolling statistics for local alerting
        self._recent_latencies: deque[float] = deque(maxlen=100)
        self._recent_confidences: deque[float] = deque(maxlen=100)
 
    def record(self, metric: InferenceMetric) -> None:
        """Record a single inference metric."""
        self.buffer.append(metric)
        self._recent_latencies.append(metric.inference_ms)
        self._recent_confidences.append(metric.confidence)
 
        # Local anomaly detection — don't wait for the cloud
        self._check_local_alerts(metric)
 
    def _check_local_alerts(self, metric: InferenceMetric) -> None:
        """Detect problems locally without needing cloud connectivity."""
        # Latency spike detection
        if len(self._recent_latencies) >= 10:
            avg = sum(self._recent_latencies) / len(self._recent_latencies)
            if metric.inference_ms > avg * 3:
                logging.warning(
                    f"Latency spike: {metric.inference_ms:.1f}ms "
                    f"(avg: {avg:.1f}ms) — possible thermal throttling"
                )
 
        # Confidence drift detection
        if len(self._recent_confidences) >= 50:
            avg_conf = (
                sum(self._recent_confidences)
                / len(self._recent_confidences)
            )
            if avg_conf < 0.5:
                logging.warning(
                    f"Low average confidence: {avg_conf:.2f} — "
                    "possible data drift or model degradation"
                )
 
        # Temperature alert
        if metric.cpu_temp and metric.cpu_temp > 80:
            logging.warning(
                f"High CPU temp: {metric.cpu_temp}°C — "
                "performance degradation likely"
            )
 
    def should_upload(self) -> bool:
        """Check if it's time to upload buffered metrics."""
        return (
            time.time() - self._last_upload > self.upload_interval
            and len(self.buffer) > 0
        )
 
    def get_upload_batch(self) -> list[dict]:
        """Get metrics for upload and clear buffer."""
        batch = [asdict(m) for m in self.buffer]
        self.buffer.clear()
        self._last_upload = time.time()
        return batch

# Sometimes the best edge strategy is... not doing edge at all.
# Here's a smart fallback pattern:
 
import aiohttp
import asyncio
 
async def hybrid_inference(
    input_data: np.ndarray,
    edge_engine: EdgeInferenceEngine,
    cloud_url: str,
    cloud_timeout_ms: float = 200,
) -> dict:
    """
    Try edge first. Fall back to cloud if edge is unhealthy.
    This gives you the best of both worlds.
    """
    # Check if edge model is healthy
    if (
        edge_engine.session is not None
        and edge_engine.avg_inference_ms < 100  # Not thermally throttled
    ):
        try:
            predictions, inference_ms = edge_engine.predict(input_data)
            return {
                "source": "edge",
                "predictions": predictions.tolist(),
                "latency_ms": inference_ms,
            }
        except Exception as e:
            logger.warning(f"Edge inference failed: {e}, falling back to cloud")
 
    # Fallback to cloud
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                cloud_url,
                json={"data": input_data.tolist()},
                timeout=aiohttp.ClientTimeout(
                    total=cloud_timeout_ms / 1000
                ),
            ) as resp:
                result = await resp.json()
                return {
                    "source": "cloud",
                    "predictions": result["predictions"],
                    "latency_ms": result.get("latency_ms", -1),
                }
    except asyncio.TimeoutError:
        logger.error("Both edge and cloud failed — returning cached result")
        return {
            "source": "cache",
            "predictions": get_cached_prediction(input_data),
            "latency_ms": 0,
        }