Edge ML Deployment

import tensorflow as tf
import numpy as np
from pathlib import Path

class MobileModelOptimizer:
    def __init__(self, model_path: str):
        self.model = tf.keras.models.load_model(model_path)
        self.tflite_model = None
        
    def optimize_for_mobile(self, 
                          quantization_type='dynamic',
                          target_spec=None,
                          representative_dataset=None):
        """Convert and optimize model for mobile deployment"""
        
        converter = tf.lite.TFLiteConverter.from_keras_model(self.model)
        
        # Set optimization flags
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        
        if quantization_type == 'int8':
            # Full integer quantization
            converter.representative_dataset = representative_dataset
            converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
            converter.inference_input_type = tf.int8
            converter.inference_output_type = tf.int8
            
        elif quantization_type == 'float16':
            # Float16 quantization
            converter.target_spec.supported_types = [tf.float16]
            
        elif quantization_type == 'dynamic':
            # Dynamic range quantization (default)
            pass
        
        # Convert model
        self.tflite_model = converter.convert()
        return self.tflite_model
    
    def benchmark_mobile_performance(self, test_data: np.ndarray, iterations=100):
        """Benchmark mobile model performance"""
        if self.tflite_model is None:
            raise ValueError("Model not optimized yet. Call optimize_for_mobile first.")
        
        # Load TFLite model
        interpreter = tf.lite.Interpreter(model_content=self.tflite_model)
        interpreter.allocate_tensors()
        
        input_details = interpreter.get_input_details()
        output_details = interpreter.get_output_details()
        
        # Warmup
        for _ in range(10):
            interpreter.set_tensor(input_details[0]['index'], test_data[:1])
            interpreter.invoke()
        
        # Benchmark
        import time
        start_time = time.time()
        
        for i in range(iterations):
            batch_idx = i % len(test_data)
            interpreter.set_tensor(input_details[0]['index'], 
                                 test_data[batch_idx:batch_idx+1])
            interpreter.invoke()
        
        avg_latency = (time.time() - start_time) / iterations * 1000  # ms
        
        # Get model size
        model_size_kb = len(self.tflite_model) / 1024
        
        return {
            'avg_latency_ms': avg_latency,
            'model_size_kb': model_size_kb,
            'throughput_qps': 1000 / avg_latency if avg_latency > 0 else 0,
            'input_shape': input_details[0]['shape'],
            'input_dtype': input_details[0]['dtype'].__name__
        }
    
    def save_mobile_model(self, output_path: str):
        """Save optimized model for mobile deployment"""
        if self.tflite_model is None:
            raise ValueError("Model not optimized yet.")
        
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, 'wb') as f:
            f.write(self.tflite_model)
        
        print(f"Optimized model saved to {output_path}")
        print(f"Model size: {len(self.tflite_model) / 1024:.2f} KB")

# Example usage
def create_representative_dataset():
    """Create representative dataset for calibration"""
    # Replace with your actual data
    for _ in range(100):
        yield [np.random.random((1, 224, 224, 3)).astype(np.float32)]

# Initialize optimizer
optimizer = MobileModelOptimizer('path/to/your/model.h5')

# Optimize for different deployment scenarios
models = {}

# 1. Dynamic quantization (fastest conversion)
models['dynamic'] = optimizer.optimize_for_mobile('dynamic')

# 2. Float16 quantization (good compression)
models['float16'] = optimizer.optimize_for_mobile('float16')

# 3. INT8 quantization (best compression)
models['int8'] = optimizer.optimize_for_mobile(
    'int8', 
    representative_dataset=create_representative_dataset()
)

# Benchmark performance
test_data = np.random.random((100, 224, 224, 3)).astype(np.float32)

for model_type, model_data in models.items():
    optimizer.tflite_model = model_data
    metrics = optimizer.benchmark_mobile_performance(test_data)
    print(f"\n{model_type.upper()} Model Performance:")
    print(f"  Latency: {metrics['avg_latency_ms']:.2f} ms")
    print(f"  Size: {metrics['model_size_kb']:.2f} KB")
    print(f"  Throughput: {metrics['throughput_qps']:.2f} QPS")
    
    # Save model
    optimizer.save_mobile_model(f'models/mobile_{model_type}.tflite')

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from typing import Dict, List, Tuple

class TensorRTOptimizer:
    def __init__(self, onnx_model_path: str):
        self.onnx_path = onnx_model_path
        self.engine = None
        self.context = None
        
        # TensorRT logger
        self.logger = trt.Logger(trt.Logger.WARNING)
        
    def build_engine(self, 
                    precision: str = 'fp16',
                    max_batch_size: int = 1,
                    max_workspace_size: int = 1 << 30):  # 1GB
        """Build optimized TensorRT engine"""
        
        builder = trt.Builder(self.logger)
        config = builder.create_builder_config()
        
        # Set memory pool limit
        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, max_workspace_size)
        
        # Parse ONNX model
        network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
        parser = trt.OnnxParser(network, self.logger)
        
        with open(self.onnx_path, 'rb') as model:
            if not parser.parse(model.read()):
                print("Failed to parse ONNX model")
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None
        
        # Set precision
        if precision == 'fp16':
            config.set_flag(trt.BuilderFlag.FP16)
        elif precision == 'int8':
            config.set_flag(trt.BuilderFlag.INT8)
            # Would need calibration dataset for INT8
        
        # Build engine
        serialized_engine = builder.build_serialized_network(network, config)
        if serialized_engine is None:
            print("Failed to build TensorRT 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()
        
        return self.engine
    
    def allocate_buffers(self) -> Tuple[List, List, List]:
        """Allocate GPU and CPU buffers"""
        inputs, outputs, bindings = [], [], []
        
        for i in range(self.engine.num_io_tensors):
            tensor_name = self.engine.get_tensor_name(i)
            dtype = trt.nptype(self.engine.get_tensor_dtype(tensor_name))
            shape = self.context.get_tensor_shape(tensor_name)
            size = trt.volume(shape)
            
            # Allocate host and device buffers
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            bindings.append(int(device_mem))
            
            if self.engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT:
                inputs.append({'host': host_mem, 'device': device_mem, 'shape': shape})
            else:
                outputs.append({'host': host_mem, 'device': device_mem, 'shape': shape})
        
        return inputs, outputs, bindings
    
    def inference(self, input_data: np.ndarray) -> np.ndarray:
        """Run inference with TensorRT engine"""
        if self.engine is None or self.context is None:
            raise RuntimeError("Engine not built. Call build_engine first.")
        
        inputs, outputs, bindings = self.allocate_buffers()
        stream = cuda.Stream()
        
        # Copy input data to GPU
        np.copyto(inputs[0]['host'], input_data.ravel())
        cuda.memcpy_htod_async(inputs[0]['device'], inputs[0]['host'], stream)
        
        # Run inference
        self.context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
        
        # Copy output from GPU
        cuda.memcpy_dtoh_async(outputs[0]['host'], outputs[0]['device'], stream)
        stream.synchronize()
        
        # Reshape output
        output_shape = outputs[0]['shape']
        return outputs[0]['host'].reshape(output_shape)
    
    def benchmark_performance(self, input_shape: Tuple[int, ...], 
                            iterations: int = 100) -> Dict[str, float]:
        """Benchmark TensorRT engine performance"""
        dummy_input = np.random.random(input_shape).astype(np.float32)
        
        # Warmup
        for _ in range(10):
            self.inference(dummy_input)
        
        # Benchmark
        import time
        start_time = time.time()
        
        for _ in range(iterations):
            self.inference(dummy_input)
        
        total_time = time.time() - start_time
        avg_latency = (total_time / iterations) * 1000  # ms
        throughput = iterations / total_time  # QPS
        
        return {
            'avg_latency_ms': avg_latency,
            'throughput_qps': throughput,
            'total_time_s': total_time
        }
    
    def save_engine(self, engine_path: str):
        """Save TensorRT engine to disk"""
        if self.engine is None:
            raise RuntimeError("No engine to save")
        
        with open(engine_path, 'wb') as f:
            f.write(self.engine.serialize())
        print(f"TensorRT engine saved to {engine_path}")

# Example usage for Jetson Nano deployment
def deploy_on_jetson():
    """Deploy optimized model on NVIDIA Jetson"""
    
    # Initialize optimizer
    optimizer = TensorRTOptimizer('model.onnx')
    
    # Build engine for Jetson Nano (limited GPU memory)
    engine = optimizer.build_engine(
        precision='fp16',  # FP16 for Jetson compatibility
        max_batch_size=1,
        max_workspace_size=512 << 20  # 512MB for Jetson Nano
    )
    
    if engine:
        # Benchmark performance
        metrics = optimizer.benchmark_performance((1, 3, 224, 224))
        print(f"Jetson Performance:")
        print(f"  Latency: {metrics['avg_latency_ms']:.2f} ms")
        print(f"  Throughput: {metrics['throughput_qps']:.2f} QPS")
        
        # Save for deployment
        optimizer.save_engine('jetson_model.trt')
        
        return optimizer
    
    return None

# Deploy the model
jetson_optimizer = deploy_on_jetson()

// Arduino/ESP32 deployment with TensorFlow Lite Micro
#include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "model_data.h"  // Generated from .tflite file

class IoTMLInference {
private:
    tflite::MicroErrorReporter micro_error_reporter;
    const tflite::Model* model;
    tflite::MicroInterpreter* interpreter;
    TfLiteTensor* input;
    TfLiteTensor* output;
    
    // Memory arena for model and intermediate arrays
    static constexpr int kTensorArenaSize = 60 * 1024;  // 60KB for ESP32
    alignas(16) uint8_t tensor_arena[kTensorArenaSize];
    
public:
    bool initialize() {
        // Load model from flash memory
        model = tflite::GetModel(model_data);
        if (model->version() != TFLITE_SCHEMA_VERSION) {
            Serial.println("Model schema version mismatch!");
            return false;
        }
        
        // Create micro op resolver with only needed operations
        static tflite::MicroMutableOpResolver<10> resolver;
        resolver.AddConv2D();
        resolver.AddDepthwiseConv2D();
        resolver.AddReshape();
        resolver.AddSoftmax();
        resolver.AddQuantize();
        resolver.AddDequantize();
        resolver.AddAdd();
        resolver.AddMul();
        resolver.AddPad();
        resolver.AddRelu();
        
        // Create interpreter
        static tflite::MicroInterpreter static_interpreter(
            model, resolver, tensor_arena, kTensorArenaSize, 
            &micro_error_reporter);
        interpreter = &static_interpreter;
        
        // Allocate memory for model tensors
        TfLiteStatus allocate_status = interpreter->AllocateTensors();
        if (allocate_status != kTfLiteOk) {
            Serial.println("AllocateTensors() failed");
            return false;
        }
        
        // Get input and output tensors
        input = interpreter->input(0);
        output = interpreter->output(0);
        
        Serial.println("IoT ML model initialized successfully");
        Serial.printf("Input shape: [%d, %d, %d, %d]\n", 
                     input->dims->data[0], input->dims->data[1], 
                     input->dims->data[2], input->dims->data[3]);
        Serial.printf("Memory usage: %d bytes\n", 
                     interpreter->arena_used_bytes());
        
        return true;
    }
    
    float* predict(const float* sensor_data, size_t data_size) {
        // Copy sensor data to input tensor
        if (input->type == kTfLiteFloat32) {
            float* input_data = input->data.f;
            for (size_t i = 0; i < data_size && i < input->bytes / sizeof(float); i++) {
                input_data[i] = sensor_data[i];
            }
        } else if (input->type == kTfLiteInt8) {
            // Handle quantized input
            int8_t* input_data = input->data.int8;
            float scale = input->params.scale;
            int32_t zero_point = input->params.zero_point;
            
            for (size_t i = 0; i < data_size && i < input->bytes; i++) {
                int32_t quantized = (sensor_data[i] / scale) + zero_point;
                quantized = max(-128, min(127, quantized));
                input_data[i] = static_cast<int8_t>(quantized);
            }
        }
        
        // Run inference
        unsigned long start_time = micros();
        TfLiteStatus invoke_status = interpreter->Invoke();
        unsigned long inference_time = micros() - start_time;
        
        if (invoke_status != kTfLiteOk) {
            Serial.println("Invoke failed");
            return nullptr;
        }
        
        Serial.printf("Inference time: %lu microseconds\n", inference_time);
        
        // Return output data
        if (output->type == kTfLiteFloat32) {
            return output->data.f;
        } else if (output->type == kTfLiteInt8) {
            // Convert quantized output to float
            static float output_float[10];  // Adjust size as needed
            int8_t* output_data = output->data.int8;
            float scale = output->params.scale;
            int32_t zero_point = output->params.zero_point;
            
            size_t output_size = output->bytes;
            for (size_t i = 0; i < output_size && i < 10; i++) {
                output_float[i] = scale * (output_data[i] - zero_point);
            }
            return output_float;
        }
        
        return nullptr;
    }
    
    void printMemoryUsage() {
        Serial.printf("Arena used bytes: %d\n", interpreter->arena_used_bytes());
        Serial.printf("Free heap: %d bytes\n", ESP.getFreeHeap());
    }
};

// Global instance
IoTMLInference ml_inference;

void setup() {
    Serial.begin(115200);
    
    // Initialize ML inference
    if (!ml_inference.initialize()) {
        Serial.println("Failed to initialize ML model!");
        while(1);
    }
    
    Serial.println("IoT ML device ready!");
}

void loop() {
    // Simulate sensor readings
    float sensor_data[28*28] = {0};  // Example for 28x28 input
    
    // Fill with actual sensor data
    for (int i = 0; i < 28*28; i++) {
        sensor_data[i] = random(0, 100) / 100.0;  // Normalized random data
    }
    
    // Run inference
    float* results = ml_inference.predict(sensor_data, 28*28);
    
    if (results) {
        // Process results (example: classification)
        float max_confidence = 0;
        int predicted_class = 0;
        
        for (int i = 0; i < 10; i++) {  # Assuming 10 classes
            if (results[i] > max_confidence) {
                max_confidence = results[i];
                predicted_class = i;
            }
        }
        
        Serial.printf("Predicted class: %d, Confidence: %.3f\n", 
                     predicted_class, max_confidence);
    }
    
    // Print memory usage periodically
    static unsigned long last_memory_check = 0;
    if (millis() - last_memory_check > 10000) {  // Every 10 seconds
        ml_inference.printMemoryUsage();
        last_memory_check = millis();
    }
    
    delay(1000);  // Run inference every second
}