TensorFlow 源码拆解

10 · TensorFlow Lite:把大 runtime 压缩成可部署解释器

FlatBuffer、Interpreter、Subgraph、delegate 与量化部署

Lite 的目标不同

TensorFlow Lite 不试图把完整 TensorFlow runtime 原样搬到手机或嵌入式设备。转换器把训练/服务侧模型变成 FlatBuffer,运行侧加载模型、构建 tensor arena、准备 subgraph,再通过 interpreter 执行有限的 builtin/custom op 集合。

SavedModel / ConcreteFunction
  → TFLite converter
  → FlatBuffer Model
  → InterpreterBuilder
  → Subgraph: tensors + nodes + execution plan
  → AllocateTensors
  → Invoke

Interpreter 的三阶段

  1. BuildInterpreterBuilder 读取 operator code 与 registration,建立模型到 runtime implementation 的映射。
  2. PrepareAllocateTensors 计算 tensor 生命周期、arena offsets、动态 shape 与 scratch buffer。
  3. Invoke:按 execution plan 调用每个 TfLiteRegistrationinit/prepare/invoke,必要时经 delegate 替换一段子图。

tensorflow/lite/core/interpreter.h 的公开 API 围绕 inputs、outputs、tensors、resize、allocate、invoke 和 delegate 修改图。与完整 runtime 相比,它把资源限制、内存复用和 ABI 稳定性放在更靠前的位置。

Subgraph 与 delegate

一个 Interpreter 可以有多个 subgraph;signature runner 用名字把输入输出映射到目标 subgraph。Delegate 先检查哪些节点可支持,再把连续节点分组为 delegate kernel,替换 execution plan。支持不完整时,剩余节点仍由 builtin interpreter 执行,因此性能和数据搬运需要观察分区边界。

量化是合同变化

int8/int16/float16 等量化路径不只是把 dtype 改小,还会引入 scale、zero point、per-axis quantization 和 kernel 实现约束。转换器、FlatBuffer schema、interpreter tensor metadata、delegate 和硬件指令必须使用同一套量化参数,否则数值会漂移或无法 prepare。

源码锚点

  • tensorflow/lite/core/interpreter.h:Interpreter 公共生命周期。
  • tensorflow/lite/core/interpreter_builder.h:模型加载、registration mapping。
  • tensorflow/lite/core/subgraph.h / subgraph.cc:tensor、node、execution plan。
  • tensorflow/lite/kernels/:builtin op registration 与实现。
  • tensorflow/lite/delegates/:GPU、NNAPI、XNNPACK、Core ML 等 delegate。
  • tensorflow/lite/schema/schema.fbs:FlatBuffer 模型 schema。
  • tensorflow/lite/python/lite.py:Python converter 入口。

一个部署侧排查顺序

先验证 FlatBuffer 能否被 builder 载入,再验证 AllocateTensors,最后看 Invoke。如果 allocate 失败,多半是 shape/arena/op prepare;如果 invoke 失败,才进入具体 kernel 或 delegate;如果结果慢,则检查 delegate 分区、fallback 和 copy。

On this page