TensorFlow Lite端侧推理框架

import tensorflow as tf # 将SavedModel转换为TFLiteconverter = tf.lite.TFLiteConverter.from_saved_model("model")tflite_model = converter.convert()with open("model.tflite", "wb") as f:    f.write(tflite_model)

converter = tf.lite.TFLiteConverter.from_saved_model("model")converter.optimizations = [tf.lite.Optimize.DEFAULT] def representative_dataset():    for _ in range(100):        data = np.random.rand(1, 224, 224, 3).astype(np.float32)        yield [data] converter.representative_dataset = representative_datasettflite_model = converter.convert()

# Android Java代码示例Interpreter interpreter = new Interpreter(modelBuffer);Interpreter.Options options = new Interpreter.Options();options.setNumThreads(4);options.addDelegate(new GpuDelegate());interpreter = new Interpreter(modelBuffer, options);

# 命令行运行Benchmarktaskset 0 ./benchmark_model \  --graph=model.tflite \  --use_gpu=true \  --num_threads=4 \  --num_runs=100

import numpy as npfrom dataclasses import dataclass, fieldfrom typing import List, Dict, Optional, Callableimport time @dataclassclass TensorInfo:    name: str    shape: list    dtype: str @dataclassclass TFLiteModel:    tensors: List[TensorInfo] = field(default_factory=list)    operators: List[dict] = field(default_factory=list)    quantization_params: Dict[str, dict] = field(default_factory=dict)    model_size: float = 0.0 class TFLiteConverter:    """模拟TFLite Converter：将TF模型概念转换为TFLite格式"""     SUPPORTED_OPS = {        "Conv2D", "DepthwiseConv2D", "FullyConnected",        "ReLU", "ReLU6", "Sigmoid", "Tanh",        "Reshape", "Concatenation", "Add", "Mul",        "AveragePool2D", "MaxPool2D", "Softmax",        "BatchNormalization", "Pad"    }     def __init__(self):        self.target_spec_ops = self.SUPPORTED_OPS.copy()        self.optimizations = []        self.representative_dataset = None     def convert(self, tf_model_graph: dict) -> TFLiteModel:        tflite_model = TFLiteModel()        unsupported = []         for op in tf_model_graph.get("operators", []):            op_type = op["type"]            if op_type in self.target_spec_ops:                tflite_model.operators.append(op)            elif op_type == "BatchNorm":                fused_op = self._fuse_batchnorm(op)                tflite_model.operators.append(fused_op)            else:                unsupported.append(op_type)         if unsupported and "SELECT_TF_OPS" not in self.target_spec_ops:            raise ValueError(f"Unsupported ops: {unsupported}")         tflite_model.tensors = [            TensorInfo(t["name"], t["shape"], t["dtype"])            for t in tf_model_graph.get("tensors", [])        ]         if "DEFAULT" in self.optimizations and self.representative_dataset:            tflite_model = self._apply_int8_quantization(tflite_model)        elif "FLOAT16" in self.optimizations:            tflite_model = self._apply_float16_quantization(tflite_model)         return tflite_model     def _fuse_batchnorm(self, bn_op: dict) -> dict:        conv_op = bn_op.get("fused_conv", {})        if conv_op:            conv_op["type"] = "Conv2D"            conv_op["fused_activation"] = "ReLU"            return conv_op        return bn_op     def _apply_int8_quantization(self, model: TFLiteModel) -> TFLiteModel:        for sample in self.representative_dataset():            for tensor in model.tensors:                if tensor.dtype == "float32":                    arr = np.array(sample, dtype=np.float32)                    scale = max(abs(arr.max()), abs(arr.min())) / 127.0                    zero_point = 0                    model.quantization_params[tensor.name] = {                        "scale": scale, "zero_point": zero_point,                        "quantized_type": "int8"                    }            break         for tensor in model.tensors:            if tensor.dtype == "float32":                tensor.dtype = "int8"         model.model_size *= 0.25        return model     def _apply_float16_quantization(self, model: TFLiteModel) -> TFLiteModel:        for tensor in model.tensors:            if tensor.dtype == "float32":                tensor.dtype = "float16"                model.quantization_params[tensor.name] = {                    "quantized_type": "float16"                }        model.model_size *= 0.5        return model  class TFLiteInterpreter:    """模拟TFLite Interpreter"""     def __init__(self, model: TFLiteModel):        self.model = model        self.delegates = []        self.num_threads = 1        self.input_data = None        self.output_data = None     def allocate_tensors(self):        pass     def set_input(self, tensor_index: int, data: np.ndarray):        self.input_data = data     def invoke(self) -> float:        start = time.perf_counter()        for op in self.model.operators:            self._execute_op(op)        elapsed = (time.perf_counter() - start) * 1000        return elapsed     def _execute_op(self, op: dict):        if self.delegates:            for delegate in self.delegates:                if delegate.can_execute(op):                    delegate.execute(op)                    return        time.sleep(0.001)     def get_output(self, tensor_index: int) -> np.ndarray:        return self.output_data if self.output_data is not None else np.array([0.0])  class Delegate:    """模拟TFLite Delegate硬件加速"""     def __init__(self, name: str, supported_ops: set):        self.name = name        self.supported_ops = supported_ops     def can_execute(self, op: dict) -> bool:        return op["type"] in self.supported_ops     def execute(self, op: dict):        time.sleep(0.0005)  class TFLiteBenchmark:    """模拟TFLite Benchmark Tool"""     @staticmethod    def run_benchmark(model: TFLiteModel, delegate: Optional[Delegate] = None,                      num_runs: int = 100) -> dict:        interpreter = TFLiteInterpreter(model)        if delegate:            interpreter.delegates = [delegate]         latencies = []        for _ in range(num_runs):            latency = interpreter.invoke()            latencies.append(latency)         latencies.sort()        return {            "avg_ms": np.mean(latencies),            "p50_ms": latencies[len(latencies) // 2],            "p99_ms": latencies[int(len(latencies) * 0.99)],            "min_ms": latencies[0],            "max_ms": latencies[-1],            "model_size_mb": model.model_size,            "delegate": delegate.name if delegate else "CPU"        }  if __name__ == "__main__":    # 构建模拟TF模型计算图    tf_graph = {        "tensors": [            {"name": "input", "shape": [1, 224, 224, 3], "dtype": "float32"},            {"name": "conv1_weight", "shape": [3, 3, 3, 32], "dtype": "float32"},            {"name": "conv1_output", "shape": [1, 224, 224, 32], "dtype": "float32"},            {"name": "fc_weight", "shape": [7, 1024], "dtype": "float32"},            {"name": "output", "shape": [1, 7], "dtype": "float32"},        ],        "operators": [            {"type": "Conv2D", "inputs": ["input", "conv1_weight"], "outputs": ["conv1_output"]},            {"type": "ReLU6", "inputs": ["conv1_output"], "outputs": ["conv1_output"]},            {"type": "FullyConnected", "inputs": ["conv1_output", "fc_weight"], "outputs": ["output"]},            {"type": "Softmax", "inputs": ["output"], "outputs": ["output"]},        ]    }     # 1. 基础转换（无量化）    converter = TFLiteConverter()    base_model = converter.convert(tf_graph)    base_model.model_size = 12.0    print(f"Base model ops: {[op['type'] for op in base_model.operators]}")    print(f"Base model tensors: {len(base_model.tensors)}")     # 2. PTQ INT8量化    def representative_data_gen():        for _ in range(10):            yield np.random.randn(1, 224, 224, 3).astype(np.float32) * 0.5     converter_int8 = TFLiteConverter()    converter_int8.optimizations = ["DEFAULT"]    converter_int8.representative_dataset = representative_data_gen    int8_model = converter_int8.convert(tf_graph)    int8_model.model_size = 12.0    print(f"\nINT8 model dtypes: {[t.dtype for t in int8_model.tensors]}")    print(f"INT8 model size: {int8_model.model_size:.1f}MB (4x compression)")     # 3. Float16量化    converter_fp16 = TFLiteConverter()    converter_fp16.optimizations = ["FLOAT16"]    fp16_model = converter_fp16.convert(tf_graph)    fp16_model.model_size = 12.0    print(f"\nFP16 model dtypes: {[t.dtype for t in fp16_model.tensors]}")    print(f"FP16 model size: {fp16_model.model_size:.1f}MB (2x compression)")     # 4. Benchmark对比    print("\n=== Benchmark Results ===")    gpu_delegate = Delegate("GPU", {"Conv2D", "FullyConnected", "ReLU6"})     results = {        "FP32 CPU": TFLiteBenchmark.run_benchmark(base_model),        "INT8 CPU": TFLiteBenchmark.run_benchmark(int8_model),        "FP16 GPU": TFLiteBenchmark.run_benchmark(fp16_model, gpu_delegate),    }     for name, result in results.items():        print(f"{name}: avg={result['avg_ms']:.2f}ms, p99={result['p99_ms']:.2f}ms, "              f"size={result['model_size_mb']:.1f}MB")     # 5. Delegate覆盖率分析    gpu_ops = [op for op in base_model.operators if gpu_delegate.can_execute(op)]    cpu_fallback = [op for op in base_model.operators if not gpu_delegate.can_execute(op)]    print(f"\nGPU Delegate coverage: {len(gpu_ops)}/{len(base_model.operators)} ops")    print(f"CPU fallback ops: {[op['type'] for op in cpu_fallback]}")