Edge AI & On-Device Inference 2026: Implementation Guide for Developers

import os
import torch
import torch.nn as nn
from executorch.exir import to_edge
from executorch.exir.backend.backend_api import to_backend

class MobileNetV3(nn.Module):
    """Lightweight model for edge deployment"""

    def __init__(self, num_classes=1000):
        super().__init__()
        # MobileNetV3-Small architecture
        self.features = nn.Sequential(
            # First conv layer
            nn.Conv2d(3, 16, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(16),
            nn.Hardswish(inplace=True),

            # MobileNetV3 blocks (simplified for brevity)
            self._make_block(16, 16, 3, 2, True, "RE", 1),
            self._make_block(16, 24, 3, 2, False, "RE", 4.5),
            self._make_block(24, 24, 3, 1, False, "RE", 3.67),
            self._make_block(24, 40, 5, 2, True, "HS", 4),
            self._make_block(40, 40, 5, 1, True, "HS", 6),
            self._make_block(40, 48, 5, 1, True, "HS", 3),
            self._make_block(48, 96, 5, 2, True, "HS", 6),
            self._make_block(96, 96, 5, 1, True, "HS", 6),

            # Final conv
            nn.Conv2d(96, 576, 1, bias=False),
            nn.BatchNorm2d(576),
            nn.Hardswish(inplace=True),
        )

        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Linear(576, 1024),
            nn.Hardswish(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(1024, num_classes),
        )

    def _make_block(self, in_ch, out_ch, k, s, se, nl, exp):
        """Helper to create MobileNetV3 block"""
        # Simplified block creation
        return nn.Sequential(
            nn.Conv2d(in_ch, int(in_ch * exp), 1, bias=False),
            nn.BatchNorm2d(int(in_ch * exp)),
            nn.ReLU(inplace=True) if nl == "RE" else nn.Hardswish(inplace=True),
            nn.Conv2d(int(in_ch * exp), int(in_ch * exp), k, s, k//2, groups=int(in_ch * exp), bias=False),
            nn.BatchNorm2d(int(in_ch * exp)),
            nn.ReLU(inplace=True) if nl == "RE" else nn.Hardswish(inplace=True),
            nn.Conv2d(int(in_ch * exp), out_ch, 1, bias=False),
            nn.BatchNorm2d(out_ch),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

# Step 1: Export PyTorch model to ExecuTorch
print("Step 1: Creating and preparing model...")
model = MobileNetV3(num_classes=10)
model.eval()

# Load pretrained weights if available
# model.load_state_dict(torch.load('mobilenetv3_weights.pth'))

example_input = torch.randn(1, 3, 224, 224)

# Trace the model
print("Step 2: Exporting to ExecuTorch edge dialect...")
try:
    # Export to torch.export format first
    exported_program = torch.export.export(model, (example_input,))

    # Convert to edge dialect
    edge_program = to_edge(exported_program)

    print(f"Edge program created successfully")
    print(f"Operations in graph: {len(edge_program.graph_module.graph.nodes)}")

except Exception as e:
    print(f"Export failed: {e}")
    raise

# Step 2: Apply quantization for size/performance
print("\nStep 3: Applying INT8 quantization...")
from executorch.exir.passes import QuantizationPass

try:
    quantized_program = edge_program.transform(
        QuantizationPass(
            quantization_config={
                "weight_dtype": "int8",
                "activation_dtype": "int8",
                "per_channel": True,  # Better accuracy than per-tensor
            }
        )
    )
    print("Quantization applied successfully")
except Exception as e:
    print(f"Quantization not available, continuing with FP32: {e}")
    quantized_program = edge_program

# Step 3: Target-specific optimization (e.g., ARM Cortex)
print("\nStep 4: Applying target-specific optimizations...")
from executorch.backends.arm import ArmBackend

try:
    arm_optimized = to_backend(
        "ArmBackend",
        quantized_program,
        compile_specs={
            "target": "cortex-a72",  # Raspberry Pi 4/5 CPU
            "enable_neon": True,     # ARM NEON SIMD instructions
            "thread_count": 4,        # Utilize all cores
            "optimize_for": "latency"  # vs "throughput" or "balanced"
        }
    )
    print("ARM optimizations applied")
except Exception as e:
    print(f"ARM backend not available: {e}")
    arm_optimized = quantized_program

# Step 4: Save for deployment
print("\nStep 5: Saving optimized model...")
output_path = "model_edge.pte"
with open(output_path, "wb") as f:
    arm_optimized.save(f)

# Analyze model size
model_size_kb = os.path.getsize(output_path) / 1024
original_size_kb = sum(p.numel() * p.element_size() for p in model.parameters()) / 1024

print(f"\n{'='*60}")
print(f"Model Export Summary:")
print(f"{'='*60}")
print(f"Original model size: {original_size_kb:.2f} KB")
print(f"Optimized model size: {model_size_kb:.2f} KB")
print(f"Compression ratio: {original_size_kb / model_size_kb:.2f}x")
print(f"Output file: {output_path}")
print(f"{'='*60}")

# Step 5: Inference example (deployment code)
print("\nExample inference code for deployment:")
print("""
# On-device inference (e.g., Raspberry Pi)
from executorch.runtime import Runtime

runtime = Runtime.get()
program = runtime.load_program('model_edge.pte')
method = program.load_method('forward')

# Prepare input
import numpy as np
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# Run inference
outputs = method.execute([input_data])
predictions = outputs[0]
""")

import os
import time
import torch
import torch.nn as nn
from torch.ao.quantization import (
    quantize_dynamic,
    get_default_qconfig_mapping,
    prepare_fx,
    convert_fx,
)
import torch.utils.data as data

class EdgeOptimizer:
    """Comprehensive optimization pipeline for edge deployment"""

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

    def optimize_for_edge(
        self,
        calibration_data: torch.utils.data.DataLoader,
        quantization_type: str = "static",
        target_backend: str = "x86"
    ) -> nn.Module:
        """
        Apply comprehensive optimization pipeline

        Args:
            calibration_data: Data loader for calibration (required for static quantization)
            quantization_type: "static", "dynamic", or "qat" (quantization-aware training)
            target_backend: "x86", "qnnpack" (ARM), or "fbgemm"

        Returns:
            Optimized model ready for edge deployment
        """
        print("Starting optimization pipeline...")
        original_size = self._get_model_size(self.model)

        # Benchmark original model
        original_latency = self._benchmark_latency(self.model, next(iter(calibration_data))[0])

        # Step 1: Pruning (remove redundant weights)
        print("\n[1/3] Applying structured pruning...")
        pruned_model = self._apply_pruning(self.model, sparsity=0.3)
        pruned_size = self._get_model_size(pruned_model)
        print(f"  Pruning: {original_size:.2f} MB → {pruned_size:.2f} MB ({original_size/pruned_size:.2f}x)")

        # Step 2: Quantization
        print(f"\n[2/3] Applying {quantization_type} quantization...")
        if quantization_type == "static":
            quantized_model = self._static_quantization(
                pruned_model,
                calibration_data,
                target_backend
            )
        elif quantization_type == "dynamic":
            quantized_model = self._dynamic_quantization(pruned_model)
        else:
            raise ValueError(f"Unknown quantization type: {quantization_type}")

        quantized_size = self._get_model_size(quantized_model)
        print(f"  Quantization: {pruned_size:.2f} MB → {quantized_size:.2f} MB ({pruned_size/quantized_size:.2f}x)")

        # Step 3: Operator fusion
        print("\n[3/3] Fusing operations...")
        fused_model = self._fuse_operations(quantized_model)

        # Final benchmarks
        final_size = self._get_model_size(fused_model)
        final_latency = self._benchmark_latency(fused_model, next(iter(calibration_data))[0])

        self.optimization_stats = {
            "original_size_mb": original_size,
            "optimized_size_mb": final_size,
            "compression_ratio": original_size / final_size,
            "quantization_type": quantization_type,
            "original_latency_ms": original_latency,
            "optimized_latency_ms": final_latency,
            "speedup": original_latency / final_latency,
        }

        print(f"\n{'='*60}")
        print("Optimization Complete:")
        print(f"  Size: {original_size:.2f} MB → {final_size:.2f} MB ({original_size/final_size:.2f}x smaller)")
        print(f"  Latency: {original_latency:.2f} ms → {final_latency:.2f} ms ({original_latency/final_latency:.2f}x faster)")
        print(f"{'='*60}\n")

        return fused_model

    def _static_quantization(
        self,
        model: nn.Module,
        calibration_data: torch.utils.data.DataLoader,
        backend: str = "x86"
    ) -> nn.Module:
        """
        Static quantization with calibration
        Best for: Maximum performance, when representative data is available
        """
        model.eval()

        # Configure quantization based on backend
        backend_map = {
            "x86": "x86",
            "fbgemm": "fbgemm",
            "qnnpack": "qnnpack",  # ARM devices
        }
        qconfig_mapping = get_default_qconfig_mapping(backend_map.get(backend, "x86"))

        # Prepare model for quantization
        example_input = next(iter(calibration_data))[0]
        prepared_model = prepare_fx(
            model,
            qconfig_mapping,
            example_inputs=(example_input,)
        )

        # Calibration pass - run representative data through model
        print("  Running calibration...")
        with torch.no_grad():
            for i, (batch, _) in enumerate(calibration_data):
                prepared_model(batch)
                if i >= 100:  # Limit calibration samples
                    break

        # Convert to quantized model
        quantized_model = convert_fx(prepared_model)

        return quantized_model

    def _dynamic_quantization(self, model: nn.Module) -> nn.Module:
        """
        Dynamic quantization (no calibration needed)
        Best for: Models with dynamic input sizes, quick optimization
        """
        return quantize_dynamic(
            model,
            {nn.Linear, nn.LSTM, nn.GRU, nn.Conv2d},  # Layers to quantize
            dtype=torch.qint8
        )

    def _apply_pruning(self, model: nn.Module, sparsity: float) -> nn.Module:
        """
        Apply structured pruning to reduce model size

        Args:
            sparsity: Fraction of weights to prune (0.0 to 1.0)
        """
        import torch.nn.utils.prune as prune

        # Create a copy to avoid modifying original
        pruned_model = model

        for name, module in pruned_model.named_modules():
            # Prune convolutional and linear layers
            if isinstance(module, nn.Conv2d):
                prune.ln_structured(
                    module,
                    name='weight',
                    amount=sparsity,
                    n=2,  # L2 norm
                    dim=0  # Prune output channels
                )
                prune.remove(module, 'weight')
            elif isinstance(module, nn.Linear):
                prune.l1_unstructured(module, name='weight', amount=sparsity)
                prune.remove(module, 'weight')

        return pruned_model

    def _fuse_operations(self, model: nn.Module) -> nn.Module:
        """
        Fuse consecutive operations for efficiency
        Common fusions: Conv2d + BatchNorm + ReLU
        """
        from torch.ao.quantization import fuse_modules

        # For quantized models, fusion is already applied during quantization
        if hasattr(model, 'qconfig'):
            return model

        # For non-quantized models, manually fuse modules
        # This is model-specific - adjust based on your architecture
        try:
            # Example fusion patterns for MobileNet-style architectures
            fused = fuse_modules(
                model,
                [
                    ['features.0', 'features.1', 'features.2'],  # Conv + BN + Activation
                    # Add more fusion patterns as needed
                ],
                inplace=False
            )
            return fused
        except Exception as e:
            print(f"  Fusion failed: {e}")
            return model  # Return original if fusion fails

    def _get_model_size(self, model: nn.Module) -> float:
        """Calculate model size in MB"""
        temp_path = "temp_model.pth"
        torch.save(model.state_dict(), temp_path)
        size_mb = os.path.getsize(temp_path) / (1024 * 1024)
        os.remove(temp_path)
        return size_mb

    def _benchmark_latency(self, model: nn.Module, sample_input: torch.Tensor, runs: int = 100) -> float:
        """Benchmark inference latency in milliseconds"""
        model.eval()

        # Warmup
        with torch.no_grad():
            for _ in range(10):
                _ = model(sample_input)

        # Benchmark
        start = time.time()
        with torch.no_grad():
            for _ in range(runs):
                _ = model(sample_input)

        latency_ms = (time.time() - start) / runs * 1000
        return latency_ms

# Usage example
if __name__ == "__main__":
    # Create sample model and data
    model = MobileNetV3(num_classes=10)

    # Create calibration dataset
    calibration_dataset = data.TensorDataset(
        torch.randn(1000, 3, 224, 224),
        torch.randint(0, 10, (1000,))
    )
    calibration_loader = data.DataLoader(calibration_dataset, batch_size=32)

    # Optimize
    optimizer = EdgeOptimizer(model)
    optimized_model = optimizer.optimize_for_edge(
        calibration_loader,
        quantization_type="static",
        target_backend="qnnpack"  # Use "qnnpack" for ARM devices
    )

    # Save optimized model
    torch.save(optimized_model.state_dict(), "optimized_model.pth")

    print("\nOptimization statistics:")
    for key, value in optimizer.optimization_stats.items():
        print(f"  {key}: {value}")

import asyncio
import torch
import torch.nn as nn
from typing import Tuple, Optional, Dict
import time
import io

class SplitInferenceModel:
    """
    Split inference between edge device and cloud
    Optimizes for latency, bandwidth, and privacy
    """

    def __init__(
        self,
        edge_layers: nn.Module,
        cloud_layers: Optional[nn.Module],
        split_point: int,
        cloud_endpoint: str,
        device: str = "cpu"
    ):
        """
        Initialize split inference system

        Args:
            edge_layers: Model layers to run on edge device
            cloud_layers: Model layers to run in cloud (None if edge-only)
            split_point: Layer index where model is split
            cloud_endpoint: API endpoint for cloud inference
            device: "cpu", "cuda", or "mps"
        """
        self.edge_model = edge_layers.to(device)
        self.cloud_endpoint = cloud_endpoint
        self.split_point = split_point
        self.device = device

        # For cloud-side processing (if running cloud service)
        self.cloud_model = cloud_layers.to(device) if cloud_layers else None

        self.edge_model.eval()
        if self.cloud_model:
            self.cloud_model.eval()

        # Performance tracking
        self.performance_history = []

    async def infer(
        self,
        input_data: torch.Tensor,
        mode: str = "auto"
    ) -> Tuple[torch.Tensor, Dict]:
        """
        Perform inference with automatic split decision

        Args:
            input_data: Input tensor (e.g., image, text embedding)
            mode: "edge_only", "cloud_only", or "auto" (automatic decision)

        Returns:
            (predictions, metadata) where metadata includes timing and decision info
        """
        metadata = {
            "edge_latency_ms": 0,
            "cloud_latency_ms": 0,
            "bandwidth_used_kb": 0,
            "execution_mode": mode,
            "timestamp": time.time()
        }

        # Step 1: Always run edge layers
        edge_start = time.time()
        with torch.no_grad():
            intermediate = self.edge_model(input_data.to(self.device))
        metadata["edge_latency_ms"] = (time.time() - edge_start) * 1000

        # Step 2: Decide execution mode
        if mode == "auto":
            mode = self._decide_execution_mode(intermediate, metadata)
            metadata["execution_mode"] = mode

        # Step 3: Execute remaining layers
        if mode == "edge_only":
            # Full edge execution (if edge model is complete)
            predictions = intermediate
        else:
            # Cloud execution
            predictions, cloud_stats = await self._cloud_inference(intermediate)
            metadata.update(cloud_stats)

        # Track performance for future decisions
        metadata["total_latency_ms"] = metadata["edge_latency_ms"] + metadata["cloud_latency_ms"]
        self.performance_history.append(metadata)

        return predictions, metadata

    def _decide_execution_mode(
        self,
        intermediate: torch.Tensor,
        edge_metadata: Dict
    ) -> str:
        """
        Intelligent decision for edge vs cloud execution

        Decision factors:
        1. Intermediate tensor size (bandwidth cost)
        2. Edge compute time (battery/thermal)
        3. Network conditions (latency, availability)
        4. Privacy requirements
        """
        # Calculate intermediate tensor size
        intermediate_size_kb = (
            intermediate.element_size() * intermediate.nelement()
        ) / 1024

        # Factor 1: Bandwidth cost
        # Large intermediate representations favor edge processing
        is_large_intermediate = intermediate_size_kb > 100  # 100 KB threshold

        # Factor 2: Edge compute time
        # If edge processing was fast, continue on edge
        is_fast_edge = edge_metadata["edge_latency_ms"] < 50  # 50ms threshold

        # Factor 3: Network conditions (simplified - would use actual network test)
        # In production, ping cloud endpoint or use cached measurements
        is_good_network = self._check_network_quality()

        # Factor 4: Privacy (simplified - would check data sensitivity)
        # For sensitive data, prefer edge processing
        is_sensitive_data = False  # Set based on data classification

        # Decision logic
        if is_sensitive_data:
            return "edge_only"  # Privacy override

        if is_large_intermediate and is_fast_edge:
            return "edge_only"  # Avoid bandwidth cost

        if not is_good_network:
            return "edge_only"  # Network issues

        # Default to cloud for complex processing
        return "cloud_only"

    def _check_network_quality(self) -> bool:
        """
        Check if network conditions are suitable for cloud inference

        In production, this would:
        - Ping cloud endpoint
        - Check bandwidth availability
        - Measure recent request latencies
        """
        # Simplified implementation
        if len(self.performance_history) > 0:
            recent = self.performance_history[-5:]  # Last 5 requests
            avg_cloud_latency = sum(
                r.get("cloud_latency_ms", 0) for r in recent
            ) / len(recent)
            return avg_cloud_latency < 200  # 200ms threshold

        return True  # Assume good network initially

    async def _cloud_inference(
        self,
        intermediate: torch.Tensor
    ) -> Tuple[torch.Tensor, Dict]:
        """
        Send intermediate representation to cloud for processing

        In production, this would use actual HTTP client (aiohttp)
        Here we simulate both client and server for demonstration
        """
        cloud_start = time.time()

        # Serialize intermediate tensor
        buffer = io.BytesIO()
        torch.save(intermediate, buffer)
        data = buffer.getvalue()
        data_size_kb = len(data) / 1024

        # Simulate network latency (would be actual API call in production)
        await asyncio.sleep(0.05)  # 50ms simulated latency

        # Cloud-side processing (if running cloud service)
        if self.cloud_model:
            with torch.no_grad():
                # Deserialize
                buffer = io.BytesIO(data)
                intermediate_cloud = torch.load(buffer)

                # Process with cloud layers
                predictions = self.cloud_model(intermediate_cloud.to(self.device))

                # Serialize result
                result_buffer = io.BytesIO()
                torch.save(predictions, result_buffer)
                result_data = result_buffer.getvalue()
        else:
            # Placeholder for actual cloud service response
            result_data = data  # Echo for demonstration

        # Calculate metrics
        cloud_latency = (time.time() - cloud_start) * 1000

        # Deserialize result
        result_buffer = io.BytesIO(result_data)
        predictions = torch.load(result_buffer)

        metadata = {
            "cloud_latency_ms": cloud_latency,
            "bandwidth_used_kb": data_size_kb,
            "response_size_kb": len(result_data) / 1024
        }

        return predictions, metadata

# Example usage
async def example_split_inference():
    """Demonstrate split inference setup"""

    # Create a simple model and split it
    full_model = nn.Sequential(
        nn.Conv2d(3, 32, 3, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Conv2d(32, 64, 3, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Flatten(),
        nn.Linear(64 * 56 * 56, 512),
        nn.ReLU(),
        nn.Linear(512, 10),
    )

    # Split after first pooling layer (layer 2)
    split_point = 3
    edge_layers = nn.Sequential(*list(full_model.children())[:split_point])
    cloud_layers = nn.Sequential(*list(full_model.children())[split_point:])

    # Create split inference model
    split_model = SplitInferenceModel(
        edge_layers=edge_layers,
        cloud_layers=cloud_layers,
        split_point=split_point,
        cloud_endpoint="https://api.example.com/infer",
        device="cpu"
    )

    # Run inference
    input_tensor = torch.randn(1, 3, 224, 224)

    # Test different modes
    for mode in ["edge_only", "cloud_only", "auto"]:
        predictions, stats = await split_model.infer(input_tensor, mode=mode)

        print(f"\nMode: {mode}")
        print(f"  Execution: {stats['execution_mode']}")
        print(f"  Edge latency: {stats['edge_latency_ms']:.2f}ms")
        print(f"  Cloud latency: {stats['cloud_latency_ms']:.2f}ms")
        print(f"  Total latency: {stats['total_latency_ms']:.2f}ms")
        print(f"  Bandwidth used: {stats['bandwidth_used_kb']:.2f}KB")

# Run example
if __name__ == "__main__":
    asyncio.run(example_split_inference())

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from typing import Optional

class JetsonDeployer:
    """Deploy optimized models on NVIDIA Jetson platforms using TensorRT"""

    def __init__(self, onnx_model_path: str):
        """
        Initialize TensorRT deployment pipeline

        Args:
            onnx_model_path: Path to ONNX format model
        """
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.engine = None
        self.context = None
        self.onnx_path = onnx_model_path

    def build_engine(
        self,
        precision: str = "fp16",
        max_batch_size: int = 1,
        workspace_size_gb: int = 1
    ) -> Optional[trt.ICudaEngine]:
        """
        Build TensorRT engine from ONNX model

        Args:
            precision: "fp32", "fp16", or "int8"
            max_batch_size: Maximum batch size for inference
            workspace_size_gb: Workspace memory in GB

        Returns:
            TensorRT engine or None if build fails
        """
        print(f"Building TensorRT engine with {precision} precision...")

        # Create builder
        builder = trt.Builder(self.logger)
        network = builder.create_network(
            1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        )
        parser = trt.OnnxParser(network, self.logger)

        # Parse ONNX model
        print(f"Parsing ONNX model from {self.onnx_path}...")
        with open(self.onnx_path, 'rb') as model_file:
            if not parser.parse(model_file.read()):
                print("ERROR: Failed to parse ONNX model")
                for error in range(parser.num_errors):
                    print(f"  Error {error}: {parser.get_error(error)}")
                return None

        print(f"Model parsed successfully. Network has {network.num_layers} layers")

        # Configure builder
        config = builder.create_builder_config()
        config.set_memory_pool_limit(
            trt.MemoryPoolType.WORKSPACE,
            workspace_size_gb << 30  # Convert GB to bytes
        )

        # Set precision mode
        if precision == "fp16":
            if builder.platform_has_fast_fp16:
                config.set_flag(trt.BuilderFlag.FP16)
                print("FP16 mode enabled (using Tensor Cores)")
            else:
                print("WARNING: FP16 not supported on this platform, using FP32")
        elif precision == "int8":
            if builder.platform_has_fast_int8:
                config.set_flag(trt.BuilderFlag.INT8)
                print("INT8 mode enabled")
                # Note: INT8 requires calibration data (not shown for brevity)
            else:
                print("WARNING: INT8 not supported on this platform, using FP32")

        # Build engine
        print("Building engine (this may take several minutes)...")
        serialized_engine = builder.build_serialized_network(network, config)

        if serialized_engine is None:
            print("ERROR: Failed to build engine")
            return None

        # Deserialize engine
        runtime = trt.Runtime(self.logger)
        self.engine = runtime.deserialize_cuda_engine(serialized_engine)
        self.context = self.engine.create_execution_context()

        print(f"Engine built successfully")

        # Save engine for faster loading next time
        engine_path = self.onnx_path.replace('.onnx', f'_{precision}.engine')
        with open(engine_path, 'wb') as f:
            f.write(serialized_engine)
        print(f"Engine saved to {engine_path}")

        return self.engine

    def load_engine(self, engine_path: str):
        """Load a previously built engine"""
        print(f"Loading TensorRT engine from {engine_path}...")

        with open(engine_path, 'rb') as f:
            serialized_engine = f.read()

        runtime = trt.Runtime(self.logger)
        self.engine = runtime.deserialize_cuda_engine(serialized_engine)
        self.context = self.engine.create_execution_context()

        print("Engine loaded successfully")

    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """
        Run inference on Jetson with TensorRT

        Args:
            input_data: Input numpy array (e.g., preprocessed image)

        Returns:
            Output predictions as numpy array
        """
        if self.engine is None or self.context is None:
            raise RuntimeError("Engine not built or loaded")

        # Get input/output binding information
        input_binding = self.engine.get_tensor_name(0)
        output_binding = self.engine.get_tensor_name(1)

        # Allocate device memory
        d_input = cuda.mem_alloc(input_data.nbytes)

        # Determine output shape
        output_shape = self.context.get_tensor_shape(output_binding)
        output_dtype = trt.nptype(self.engine.get_tensor_dtype(output_binding))
        output_data = np.empty(output_shape, dtype=output_dtype)
        d_output = cuda.mem_alloc(output_data.nbytes)

        # Create stream for async execution
        stream = cuda.Stream()

        # Transfer input data to device
        cuda.memcpy_htod_async(d_input, input_data, stream)

        # Set tensor addresses
        self.context.set_tensor_address(input_binding, int(d_input))
        self.context.set_tensor_address(output_binding, int(d_output))

        # Execute inference
        self.context.execute_async_v3(stream_handle=stream.handle)

        # Transfer predictions back to host
        cuda.memcpy_dtoh_async(output_data, d_output, stream)

        # Synchronize stream
        stream.synchronize()

        return output_data

    def benchmark(self, input_shape: tuple, num_iterations: int = 100) -> dict:
        """
        Benchmark inference performance

        Returns:
            Dictionary with latency statistics
        """
        # Generate random input
        input_data = np.random.randn(*input_shape).astype(np.float32)

        # Warmup
        for _ in range(10):
            _ = self.infer(input_data)

        # Benchmark
        import time
        latencies = []

        for _ in range(num_iterations):
            start = time.time()
            _ = self.infer(input_data)
            latencies.append((time.time() - start) * 1000)  # Convert to ms

        return {
            "mean_latency_ms": np.mean(latencies),
            "std_latency_ms": np.std(latencies),
            "min_latency_ms": np.min(latencies),
            "max_latency_ms": np.max(latencies),
            "throughput_fps": 1000 / np.mean(latencies),
        }

# Usage example
if __name__ == "__main__":
    # First, convert PyTorch model to ONNX (prerequisite)
    """
    import torch
    model = YourModel()
    model.eval()
    dummy_input = torch.randn(1, 3, 224, 224)
    torch.onnx.export(
        model,
        dummy_input,
        "model.onnx",
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
    )
    """

    # Deploy to Jetson
    deployer = JetsonDeployer("model.onnx")

    # Build engine (first time only)
    engine = deployer.build_engine(
        precision="fp16",  # Leverage Jetson's Tensor Cores
        max_batch_size=1
    )

    # For subsequent runs, load pre-built engine:
    # deployer.load_engine("model_fp16.engine")

    # Run inference
    input_array = np.random.randn(1, 3, 224, 224).astype(np.float32)
    predictions = deployer.infer(input_array)

    print(f"Predictions shape: {predictions.shape}")
    print(f"Top prediction: {np.argmax(predictions)}")

    # Benchmark performance
    stats = deployer.benchmark((1, 3, 224, 224), num_iterations=100)
    print(f"\nPerformance Statistics:")
    for key, value in stats.items():
        print(f"  {key}: {value:.2f}")

# Convert ONNX to HEF using Hailo Dataflow Compiler
hailo parser onnx model.onnx --output model.har
hailo optimize model.har --output model_optimized.har
hailo compiler model_optimized.har --output model.hef

edgetpu_compiler model.tflite

# Model metadata for version management
model_manifest = {
    "model_id": "mobilenet_v3_small",
    "version": "1.2.0",
    "checksum": "sha256:abc123...",
    "size_bytes": 2_500_000,
    "requires_runtime_version": ">=1.0.0"
}

# Pseudo-code for adaptive power management
if battery_level < 20:
    set_cpu_governor("powersave")
    reduce_inference_frequency()
elif thermal_throttling_active:
    reduce_model_complexity()  # Switch to lighter model
else:
    set_cpu_governor("performance")