已有 .tflite 模型转 int8 量化

有一个现成的 sine_calcu.tflite 浮点模型(2380 字节,下载),需要量化为 int8 给 STM32 用。打开终端,装好 TF 2.21,写下一行熟悉的代码:

converter = tf.lite.TFLiteConverter.from_frozen_buffer(model_data, ...)

报错:AttributeError: type object 'TFLiteConverterV2' has no attribute 'from_frozen_buffer'

查了一圈发现 TF 2.16+ 重构了 TFLiteConverter 的 API。新的构造函数签名是 TFLiteConverter(funcs, trackable_obj=None),需要 concrete functions。但手里只有一个 .tflite 文件——没有 SavedModel,没有 Keras 模型,也没有原始训练代码。concrete function 从哪来?

这篇文章记录了这个问题的解法:用 flatbuffers 解析 .tflite 模型的内部结构,提取所有权重张量,在 Keras 中重建一个完全相同的模型,再走标准的 Keras → int8 量化流程。

第一步:用 flatbuffers 解析模型结构

TensorFlow 自带 TFLite 的 flatbuffers schema(tensorflow.lite.python.schema_py_generated),可以直接解析 .tflite 文件的内部结构:

from tensorflow.lite.python import schema_py_generated as _schema

with open("sine_calcu.tflite", "rb") as f:
    model_data = f.read()

model = _schema.Model.GetRootAs(model_data, 0)
sg = model.Subgraphs(0)

for i in range(sg.OperatorsLength()):
    op = sg.Operators(i)
    oc = model.OperatorCodes(op.OpcodeIndex())
    inputs = [op.Inputs(j) for j in range(op.InputsLength())]
    outputs = [op.Outputs(j) for j in range(op.OutputsLength())]
    print(f"Op [{i}]: builtin_code={oc.BuiltinCode()}, "
          f"inputs={inputs}, outputs={outputs}")

输出:

Op [0]: builtin_code=9, inputs=[0, 5, 6], outputs=[7]    ← FULLY_CONNECTED
Op [1]: builtin_code=28, inputs=[7], outputs=[8]         ← TANH
Op [2]: builtin_code=9, inputs=[8, 4, 2], outputs=[9]    ← FULLY_CONNECTED
Op [3]: builtin_code=28, inputs=[9], outputs=[10]        ← TANH
Op [4]: builtin_code=9, inputs=[10, 3, 1], outputs=[11]  ← FULLY_CONNECTED

模型结构一目了然:Dense(1→10) → Tanh → Dense(10→5) → Tanh → Dense(5→1),一个三层 MLP。每个 tensor 的名称和 shape 也都能从 flatbuffer 中读到。

第二步:提取权重张量

所有权重都存储在 flatbuffer 的 buffer 区。根据 tensor 的名称和索引关系,定位到 6 个权重张量:

tensor_data = {}
for i in range(sg.TensorsLength()):
    t = sg.Tensors(i)
    name = t.Name().decode() if t.Name() else f"tensor_{i}"
    buf = model.Buffers(t.Buffer())
    data_len = buf.DataLength()
    if data_len > 0:
        shape = [t.Shape(j) for j in range(t.ShapeLength())]
        raw_bytes = buf.DataAsNumpy()
        arr = np.frombuffer(raw_bytes, dtype=np.float32).reshape(shape).copy()
        tensor_data[name] = arr

提取结果:

sequential_1/dense_1/MatMul   : [10, 1]   ← Layer 1 kernel
sequential_1/dense_1/BiasAdd  : [10]      ← Layer 1 bias
arith.constant3               : [5, 10]   ← Layer 2 kernel
arith.constant1               : [5]       ← Layer 2 bias
arith.constant2               : [1, 5]    ← Layer 3 kernel
arith.constant                : [1]       ← Layer 3 bias

注意 TFLite 中 FC 层的权重 shape 是 [out_features, in_features],Keras 的 Dense 是 [in_features, out_features],赋值时需要 transpose。

第三步:在 Keras 中重建模型

用提取的权重构建一个完全相同的 Keras 模型:

model_keras = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='tanh', input_shape=(1,)),
    tf.keras.layers.Dense(5, activation='tanh'),
    tf.keras.layers.Dense(1),
])

# Build
_ = model_keras(np.array([[0.0]], dtype=np.float32))

# Set weights — 注意 transpose
model_keras.layers[0].set_weights([w1.T, b1])
model_keras.layers[1].set_weights([w2.T, b2])
model_keras.layers[2].set_weights([w3.T, b3])

验证 float 精度——与原 .tflite 模型对比:

x=-2.00  keras=+0.643296  tflite=+0.643297  OK
x=-1.00  keras=+0.243617  tflite=+0.243617  OK
x=+0.00  keras=+0.000132  tflite=+0.000133  OK
x=+0.50  keras=-0.001474  tflite=-0.001474  OK
x=+1.00  keras=+0.014659  tflite=+0.014659  OK
x=+2.00  keras=+0.038101  tflite=+0.038101  OK

六个测试点误差在 1e-6 级别,权重提取和重建完全正确。

第四步:Full Int8 量化

现在有了 Keras 模型,就可以走标准的 TF 量化流程:

def representative_dataset():
    for _ in range(500):
        x = np.random.uniform(-np.pi, np.pi, size=(1, 1)).astype(np.float32)
        yield [x]

converter = tf.lite.TFLiteConverter.from_keras_model(model_keras)
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.int8
converter.inference_output_type = tf.int8

tflite_quant_model = converter.convert()

这里的关键参数:

量化结果

Float model: 2380 bytes
Int8  model: 2992 bytes

int8 模型尺寸不降反升——这个模型太小了(总共才 6 个权重张量,加起来不到 300 个 float),新增的 Quantize/Dequantize 算子和量化元数据开销超过了权重本身的压缩收益。对于真实规模的模型(几 MB 以上),int8 通常会缩小 3-4 倍。

量化参数:

Scale Zero Point
Input 0.02451 0
Output 0.00341 -126

精度验证——用 int8 推理并与原 float 模型对比:

x=+0.00  int8=+0.00683  float_ref=+0.00013  error=0.0067
x=+0.50  int8=+0.00683  float_ref=-0.00147  error=0.0083
x=+1.00  int8=+0.01024  float_ref=+0.01466  error=0.0044
x=+1.50  int8=+0.04096  float_ref=+0.02848  error=0.0125
x=-1.00  int8=+0.24919  float_ref=+0.24362  error=0.0056
x=+3.14  int8=+0.06145  float_ref=+0.05494  error=0.0065

误差在 int8 量化的正常范围(~0.01 级别)。量化后模型的函数语义保持不变。

量化前后的模型文件:

这个方法能泛化吗

这个流程的核心思路是 "从 .tflite flatbuffer 反向工程出模型结构和权重",它适用于以下条件:

对于更复杂的模型(含有分支、跳跃连接、自定义算子),手动重建会变得困难。那种场景下,更好的做法是找到原始训练代码,或者用 TF 1.x 的 tf.compat.v1.lite.TFLiteConverter.from_frozen_graph 降级处理。

这次量化工作的本质是把一个不透明的二进制模型文件变成了可操作的 Keras 对象。flatbuffers 在这里扮演了"解包器"的角色——即便 TFLite 的序列化格式是 Protobuf 或其他格式,思路一样:解析序列化格式,提取数据,重建计算图。