.tflite — TFLite模型格式的实质、结构及应用

.tflite格式的实质

.tflite是机器学习模型的序列化表示，针对在内存和计算能力有限的设备上进行推理进行了优化。与将图形存储在protobuf中并需要大量资源加载的SavedModel格式不同，.tflite使用FlatBuffers — 一种无需复制即可访问数据的序列化库。

.tflite格式在设计时考虑了移动平台的特点：加载时内存消耗最小、快速启动、支持量化权重。.tflite文件不支持训练 — 仅支持推理。这是与完整TensorFlow格式的根本区别，可以显著减小运行时的大小。

根据Google Research, 2024的数据，采用INT8量化的.tflite格式模型比采用FP32格式的类似模型加载速度快60%，并且在推理期间消耗的RAM少40%。

.tflite文件的内部结构

.tflite文件由多个部分组成，按照FlatBuffers方案组织。主要部分 — Model，包含计算图（SubGraph）、运算符列表（OperatorCodes）和权重缓冲区。每个SubGraph包含张量、运算符、输入和输出索引。

OperatorCodes部分包含模型中使用的所有运算符的标识符 — Conv2D、DepthwiseConv2D、FullyConnected、Softmax等。每个运算符都有一个来自预定义BuiltinOperator列表的代码。如果模型包含列表中不存在的操作，则它们被标记为Custom，需要通过委托或自定义运算符来实现。

Buffers部分包含模型权重的二进制数据。根据量化模式，权重可以存储在FP32、FP16、INT8或带有缩放参数的量化格式中。缓冲区通过mmap机制直接映射到内存，从而消除了额外的复制。

.tflite文件的主要部分

量化和格式优化

.tflite支持三种量化模式。完全整数INT8量化 — 权重和激活用8位整数表示。转换需要代表性数据集来校准激活范围。模型大小缩小4倍。

FP16量化 — 权重转换为16位浮点数。此模式无需校准，精度损失最小。激活保持FP32格式。模型大小缩小2倍。推荐用于GPU和NPU上支持FP16的设备。

动态量化 — 权重转换为INT8，但激活以FP32计算。模型大小缩小4倍，但精度仍高于完全INT8。根据Google ML Performance Benchmarks，此模式对于NLP模型和对精度高度敏感的模型是最佳的。

从TensorFlow模型创建.tflite

要将模型转换为.tflite，使用TFLite Converter — 通过Python API tf.lite.TFLiteConverter可访问的TensorFlow组件。Converter支持三种来源：SavedModel、Keras H5和ConcreteFunction。转换Keras模型的最小示例：

import tensorflow as tf

# 从文件加载训练好的Keras模型
model = tf.keras.models.load_model("my_model.h5")

# 使用动态量化为.tflite格式
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

# 将转换后的模型写入.tflite文件
with open("model.tflite", "wb") as f:
    f.write(tflite_model)

要使用完全INT8量化进行转换，需要校准数据集。指定representative_dataset以确定激活范围：

converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
tflite_quant_model = converter.convert()

用于.tflite的工具

Google提供一套工具用于分析和优化.tflite模型。TFLite Benchmark Tool测量设备上的推理时间并显示每个运算符的统计信息。TFLite Model Inspector可视化计算图并显示张量大小。

Netron — 跨平台模型可视化工具，支持.tflite以及ONNX、Core ML和PyTorch。Netron显示图形结构、每个运算符的参数以及张量之间的连接。在转换过程中调试时很有用。

TFLite Metadata API允许向.tflite文件添加元数据：模型描述、输入数据格式、测量单位、作者信息。元数据由Task Library用于自动配置预处理和后处理。

加载和执行.tflite的示例

在iOS上使用Swift API加载.tflite模型的示例。TFLite Swift API提供Interpreter用于从bundle加载模型并执行推理。通过Core ML Delegate指定委托以实现硬件加速：

import TensorFlowLite

guard let modelPath = Bundle.main.path(
    forResource: "model", ofType: "tflite"
) else { return }

let interpreter = try Interpreter.init(modelPath: modelPath)
try interpreter.allocateTensors()

// 为模型准备输入数据
let inputData = Data.init(count: 1 * 224 * 224 * 3)
try interpreter.copy(inputData, toInputAt: 0)

// 运行模型推理
try interpreter.invoke()

// 读取输出张量数据
let outputTensor = try interpreter.output(at: 0)
let results = outputTensor.data.withUnsafeBytes {
    Array($0.bindMemory(to: Float32.self))
}