TensorFlow Lite概述：转换器、解释器、XLA和2019年路线图 - 极术社区

导读：虽然TensorFlow用户众多，但是推理框架这方面，其他家做的都很好。就从我的角度来看，市面上较为知名的如腾讯的NCNN，因其开源早，且社区一直有作者在维护，文档详细、例子充分、业内使用最为广泛，也基于Vulkan支持手机GPU。再就是小米自研的MACE，整体代码和文档质量高，不仅支持CPU、GPU（高通Adreno GPU的实现性能与SNPE不相上下），也支持Hexagon DSP，据说Snapchat用的男女变换的模型，就是用MACE作为推理引擎加载的。
在我看来，好的推理引擎，不光是性能，没有从用户出发、开箱即用、及时更新的丰富经典代码案例，那走不远的，再去谈用户群、及时的issue解答、可扩展性等特色。毕竟框架引擎是拿来给开发者用的，不是只看收藏数这些东西。在各家框架互相致敬的大背景下，TensorFlow Lite这些方面也是可圈可点。
文档，代码以2019年07月06日附近的日期为主：
日期：2019年07月06日
repo：https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite
commit：718503b075d7b
微信公众号：NeuroMem
Editor: https://github.com/ysh329

本文主要从文档的宏观角度去看TensorFlow Lite的现状，不分析代码，了解它在易用性、性能等方面的优点，说白了就是整合了下文档。目录如下：

1. 概述：特性、组成模块、部署流程

TensorFlow Lite是一个用于将TensorFlow模型部署到移动、嵌入式、物联网设备上的低延迟，轻量推理框架。其特点如下：

为不同端上优化的核心operator的解释器（Interpreter）打包成一个轻量的二进制包；
丰富的平台支持。Android和iOS设备、嵌入式Linux、微控制器设备等；
多语言API调用支持。Java, Swift, Objective-C, C++, and Python；
高性能。针对硬件精心优化的底层kernel实现，如预融合策略的激活和bias算子/底层kernel；
模型优化工具。量化模型，该策略可在不影响精度的前提下减小模型尺寸；
高效模型格式。使用轻量的FlatBuffer（跨平台，且比protobuf性能好，最初为游戏和性能相关任务而开发）；
提供常见任务的预训练模型，以方便用户快速定制化应用；
提供丰富的案例和教程教你在不同平台和设备上部署机器学习模型。

根据其官方guide，TensorFlow Lite由两部分组成：

解释器（interpreter）：在众多不同的硬件型号上如手机、Linux系统的嵌入式硬件和微控制器（microcontrollers），运行（优化过的）模型；
转换器（converter）, 将TensorFlow模型转换为解释器（interpreter）可用的高效形式，并做了优化以可减小二进制文件的尺寸（当链接所有op时，Lite的二进制包小于300KB，当只包括inceptionV3和MobileNet模型的对应op时，包小于200KB）并提高性能。

部署流程有4个步骤：

选择待部署的模型，如预训练的模型或者来自retrain的模型；
转换模型为Lite格式。若模型为自己用TF训练的，使用转换器（Convertor）转换为Lite格式；
在端上使用Lite的解释器（interpreter）支持的API，来运行Lite模型；
优化模型：使用模型优化工具，在牺牲微小的精度下，减小尺寸提供推理速度。

2. 转换器（Convertor）的优化流程

将训练好的模型部署到设备上，需要使用转换器（有Python API和命令行两种方式，官方推荐具有更好op兼容性的前者，转换时也需要注意从TensorFlow到其Lite形式的算子兼容性）。

转换器（Convertor）将TensorFlow常见的四种格式（SaveModels、freeze_graph.py生成的Frozen GraphDef、tf.keras HDF5 models、tf.Session TF-Python-API generated）的模型转换成Lite的FlatBuffer格式.tflite。整个过程如下图所示：

图上可以看到，Lite的解释器（Interpreter）使用的是压缩后的FlatBufferModel用于推理。其中，FlatBuffers是一个跨平台、高效、开源的序列化库，类似protocol buffers，但前者特点在其访问数据前不需要做解析或解压缩（parsing/unpacking）的步骤，避免了内存分配。

2.1 模型优化（model optimization）

就目前官方提供的工具来说，模型优化指的是量化和权重剪枝。

2.1.1 模型量化

先来看一组量化的benchmark数据：

测试平台：谷歌Pixel 2；
平台参数：Intel visual core + 骁龙835处理器（Adreno540、CPU为Kryo280 1.9Ghz + 2.45Ghz）；
两种量化方式：post-training quantization、quantization-aware training。

量化的优势：

支持现有的CPU平台，使用成本低；
激活的量化可以减少内存访问读和存（中间临时量）的时间和空间开销；
大多CPU和硬件加速器都支持单指令多数据的指令，有助于量化后的kernel加速。

Lite提供两种量化方式（several levels of support for quantization）：

非训练量化（Post-training: Quantizing models for CPU model size，官方推荐）：在推理过程中，将权重量化为8位并“在运行中”量化输入/激活。性能提升的收益，大于模型量化过程中带来的内存变化：

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
tflite_quant_model = converter.convert()

训练量化（During training: Quantizing models for integer-only execution.）：用于仅整数执行的量化模型获得具有更快、更小尺寸和仅整数加速器兼容的模型。目前，该方式的训练需要有“假量化”节点的模型，且只支持一小部分的网络结构。

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.inference_type = tf.lite.constants.QUANTIZED_UINT8
input_arrays = converter.get_input_arrays()
converter.quantized_input_stats = {input_arrays[0] : (0., 1.)}  # mean, std_dev
tflite_model = converter.convert()

对全整型模型，输入是无符号的uint8类型。均值和标准差反应了uint8的值在训练过程中于float类型间的映射关系。

均值也是介于0到255之间的值，映射到了浮点的0.0f，标准差的计算是255 / (float_max - float_min)。

官方也在积极开发基于这两种量化方法的新工具，以期简化量化工具的使用难度。

2.1.2 权重剪枝（Weight pruning）

若pre-optimized模型（官方提供了一些在ImageNet上预训练好的量化模型，如MobileNet等）、post-training工具都不满足您的性能要求，那还可以尝试不同的训练工具。
训练过程中的优化工具（Training time tools），是基于对模型损失函数的修改，从而使模型可以“适应”优化技术带来的变化，达到优化的目的。其实现是基于Keras的训练脚本，从一个预训练过的Keras模型来做参数微调。目前训练过程中的优化工具（Training time tools），能用的有：
权重剪枝（Weight pruning）
2.量化（Quantization）：这个仍在开发过程中

基于量级的权重剪枝，会在训练过程中将权重逐步归零，以达到稀疏性目的。稀疏模型的优势在于易于可压缩，在推理过程中跳过权重为0的计算。剪枝通过压缩模型来实现。在牺牲微小精度的前提下，我们的模型有6倍的性能提升。剪枝这项技术，也在实际应用到语音应用中，如语音识别、文本转语音以及横跨各种视觉与翻译的模型中。

用户可以在TensorFlow 1.14+的版本，通过 Keras API 以静态图（Graph）和动态图（eager execution）方式来使用这项技术。

权重剪枝，也是量化后对模型进一步优化的方法，通过对模型减去不重要的权重连接来减小计算量和参数量。官方提供了详尽的Keras剪枝教程和Python API文档，以及训练稀疏模型等高级用法的指导。此外，还有一个MNIST手写体图像识别的CNN模型剪枝代码，和IMDN情感分类的LSTM模型的剪枝代码也可以参考。

先来看两组剪枝模型前后的对比性能。第一组是基于ImageNet数据集训练的图像分类模型：

第二组是基于 WMT16 German 和 English 数据集的翻译模型，其中 news-test2013 作为训练集、 news-test2015 作为测试集：

官方给出了剪枝的技巧：

最好从一个预训练或者训练后的模型剪枝，不要一开始就剪枝；
剪枝过程中的剪枝频率不要太频繁（一般剪枝工具也会提供默认的剪枝频率），给模型一些训练的时间以恢复；
尝试从预训练模型的剪枝实验时，step从0开始一直到最终的稀疏（我这句没看懂：Try running an experiment where you prune a pre-trained model to the final sparsity with begin step 0）；
剪枝的策略也是一个超参数，模型剪枝过程中的学习率不要设置的过高或者过低；

更多关于剪枝的背景，可以看这篇文章：To prune, or not to prune: exploring the efficacy of pruning for model compression。

3. 解释器（Interpreter）的推理流程

推理过程是通过解释器（interpreter）来执行。解释器使用静态图命令（static graph ordering）和一个定制的内存分配器以确保最小的负荷、初始化和执行延迟。Lite支持常见移动和嵌入式平台，如Android、iOS和Linux。

推理的重要概念

读取模型：将.tflite模型加载到内存中，其中包含了模型的执行流图；
数据转换（Transforming Data）：将输入数据转换成模型接收的形式或排布，如resize原始图像到模型输入大小；
执行推理（Running Inference）：这一步使用API来执行模型。其中包括了创建解释器、分配张量等，具体见后文Running a Model；
解释输出（Interpreting Output）：用户取出模型推理的结果，并解读输出，如分类结果的概率。

3.1 读取模型（Loading a Model）

class FlatBufferModel {
  // 基于文件构建模型，若创建失败返回nullptr
  static std::unique_ptr<FlatBufferModel> BuildFromFile(
      const char* filename,
      ErrorReporter* error_reporter);

  // 基于预加载（pre-loaded）的flatbuffer构建模型。
  // 调用者保留缓冲区的所有权，并应保持活动状态，直到返回的对象被销毁
  // 构建失败返回nullptr
  static std::unique_ptr<FlatBufferModel> BuildFromBuffer(
      const char* buffer,
      size_t buffer_size,
      ErrorReporter* error_reporter);
};
tflite::FlatBufferModel model(path_to_model);

注意，如果TensorFlow Lite检测到存在Android NNAPI，它会自动尝试使用共享内存来存储FlatBufferModel。

3.2 运行模型（Running a Model）

该过程包含如下步骤：

基于FlatBufferModel来构建解释器（Interpreter）；
如有必要，对输入数据做如缩放等预处理操作；
填充input tensor的值；
计算推理；
读出输出结果。

这里给出解释器（interpreter）公共接口的重要部分，需要注意：

为了避免字符串比较（没看懂这句话：in order to avoid string comparisons (and any fixed dependency on string libraries)），张量的表示都用整型（integers）；
不能从并发线程访问解释器；
输入和输出张量的内存分配，必须在调整张量的大小后立即调用 AllocateTensors() 来触发；

Lite的解释器（Interpreter）加载的模型格式，必须为FlatBufferModel格式。在解释器整个生命周期内，都必须保证FlatBufferModel是可用的。

此外，单个FlatBufferModel可以被多个解释器同时使用。FlatBufferModel创建了后，才能被解释器加载（注意前后关系），且要保证解释器使用完了再销毁。

Lite的解释器C++使用方式如下：

// 加载模型
tflite::FlatBufferModel model(path_to_model);

// 构建解释器
tflite::ops::builtin::BuiltinOpResolver resolver;
std::unique_ptr<tflite::Interpreter> interpreter;
tflite::InterpreterBuilder(*model, resolver)(&interpreter);

// 如有需要，对输入tensor缩放
interpreter->AllocateTensors();

// 填充输入tensor + 预处理（略）
float* input = interpreter->typed_input_tensor<float>(0);
// Fill `input`.

// 执行推理
interpreter->Invoke();

// 取出结果 + 后处理（略）
float* output = interpreter->typed_output_tensor<float>(0);

解释器的Python接口使用方式分两种：load a model file和from model data，这里给出load a model a file的方式：

import numpy as np
import tensorflow as tf

# 读取TFLite模型并分配tensors
interpreter = tf.lite.Interpreter(model_path="converted_model.tflite")
interpreter.allocate_tensors()

# 获取输入和输出tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 用随机输入测试模型
input_shape = input_details[0]['shape']
input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)

interpreter.invoke()

# 函数`get_tensor()`返回该tensor数据的一份拷贝
# 使用`tensor()`返回该tensor的指针
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

3.3 写自定义Operator

TensorFlow Lite的operator（包括用户custom自定义和builtin已有的），都是用纯C语言接口，其中包含四个函数：

typedef struct {
  void* (*init)(TfLiteContext* context, const char* buffer, size_t length);
  void (*free)(TfLiteContext* context, void* buffer);
  TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node);
  TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node);
} TfLiteRegistration;

有关TfLiteContext和TfLiteNode的详细信息，请参阅context.h：

前者提供错误报告和对全局对象的访问功能，包括所有张量；
后者允许访问其输入和输出的实现。

当解释器（interpreter）加载模型时，它会为图中的每个节点调用一次init()。如果在图中多次使用某个operator，则会多次调用该init()函数。对于自定义operator，将提供配置缓冲区，其中包含将参数名称映射到其值的flexbuffer。内置的（builtin）operator的缓冲区为空，因为解释器已经解析了op参数。需要状态的内核实现，应该在此处初始化并将所有权转移给调用者（caller）。对于每个init()调用，都会有与之相对应的free()用来分配当初在init()过程中分配的资源。

每当输入张量调整大小时，解释器将通过图（graph）通知实现的变更。通知后，内部缓冲区的大小会做出调整，也会检查输入形状和类型的有效性，并重新计算输出形状。这一切都是通过prepare()完成的，实现时可以使用node->user_data来访问它们的状态。

每次推断运行时，解释器遍历图（graph）调invoke()，这里的状态也可用node->user_data来访问。

通过定义这四个函数和通常如下所示的全局注册函数，可以让自定义实现的operator和内置operator完全相同：

namespace tflite {
namespace ops {
namespace custom {
  TfLiteRegistration* Register_MY_CUSTOM_OP() {
    static TfLiteRegistration r = {my_custom_op::Init,
                                   my_custom_op::Free,
                                   my_custom_op::Prepare,
                                   my_custom_op::Eval};
    return &r;
  }
}  // namespace custom
}  // namespace ops
}  // namespace tflite

注册不是自动的完成的，需要显式调用Register_MY_CUSTOM_OP来实现。虽然标准的BuiltinOpResolver（可从builtin_ops这个target获得）负责内置算子的注册，但也必须在在某个自定义库中将这些自定义的算子集合起来（我理解就是需要写个头文件里面都是这些算子的注册声明）。

3.4 写自定义Kernel库

实际上，解释器会加载一个内核（Kernel）库，该内核库会被分配给对应算子。虽然默认库只包含内置内核（builtin kernels），但可以用自定义的内核库替换它。解释器使用OpResolver将算子代码和其名称转为实际代码。

class OpResolver {
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  virtual void AddOp(tflite::BuiltinOperator op, TfLiteRegistration* registration) = 0;
  virtual void AddOp(const char* op, TfLiteRegistration* registration) = 0;
};

常规的用法需要开发人员使用BuiltinOpResolver来写：

Regular usage will require the developer to use the BuiltinOpResolver and write:

tflite::ops::builtin::BuiltinOpResolver resolver;

注册自定义的算子这么写：

resolver.AddOp("MY_CUSTOM_OP", Register_MY_CUSTOM_OP());

在解析器传递给InterpreterBuilder之前（感觉文档里这句话没啥用）。若内置算子的库体积太大，可基于给定算子如仅包含给定模型的算子，来代码生成新的OpResolver。等同于 TensorFlow 的选择性注册（ selective registration），有个简化版的实现位于tools目录下。

4. XLA概述

XLA（Accelerated Linear Algebra，加速线性代数）是一种优化TensorFlow计算的编译器。 XLA（加速线性代数）是为优化TensorFlow计算的特定领域编译器。它可以提高服务器和移动平台的速度、改进内存使用、提升便携性（或者说是可移植性）。 XLA框架是目前处于开发中的实验性项目。更多信息请阅读XLA指南，咱们下面就开始讲这个指南的内容。

鼓励针对新硬件加速器的开发人员试用XLA。XLA框架是实验性的并且在积极开发中。虽然现有算子的语义不太可能发生变化，但预计会增加更多算子来覆盖重要的用例。

4.1 XLA项目的初衷

提高执行速度。编译子图以减少短期Ops的执行时间，以消除TensorFlow运行时的开销，融合流水线操作以减少内存开销，并专注于已知的张量形状以允许更积极的恒定传播（constant propagation）；
提高内存使用率。分析和安排内存使用，原则上消除了许多中间存储缓冲区；
减少对自定义Ops的依赖。通过提高自动融合的低级Ops的性能来匹配手动融合的自定义Ops的性能，从而消除了对许多自定义Ops的需求；
减少移动端的库大小。通过提前编译子图并发出可直接链接到另一个应用程序的对象/头文件对来消除TensorFlow运行时。结果可以将移动推断的占用空间减少几个数量级；
提高可移植性。为新硬件编写新的后端相对容易，很大一部分TensorFlow程序，不修改代码的情况下可以在新硬件上运行。相比为新硬件专门设计单个Ops，不需要重写针对该硬件的Ops。

4.2 XLA如何工作的

XLA的输入语言称为“HLO IR”，或HLO（高级优化器，我觉得叫做“高级优化编译器”更合适）。 HLO的语义在Operation Semantics有描述，其实HLO的表示，即编译器IR。

XLA会使用HLO中定义的图（Graph，也是“计算”），并将图（graph）编译为各种体系结构的机器指令。XLA的模块化设计，使它很容易加入一个新硬件架构的后端，目前已经支持的后端有x64和ARM64的CPU后端以及NVIDIA GPU后端，下图展示了XLA中的编译过程：

XLA有几个与目标硬件无关的优化，例如公共子表达式消除（CSE，common subexpression elimination），过程与硬件后端无关的（target-independent）算子融合，以及用于为计算而在运行时做内存分配的缓冲区分析。

硬件目标无关的步骤之后，XLA将HLO计算发送到后端去做进一步的HLO级优化，HLO级优化会针对硬件目标特定信息和需求，如XLA GPU后端的算子融合、确定如何将计算划分为流。此阶段，后端还可以将某些算子或其组合与优化的库调用进行模式匹配。

下一步是特定于目标硬件的代码生成。CPU包括XLA使用的GPU后端会使用LLVM生成低级IR、优化和代码生成。硬件后端以有效方式发出了表示XLA HLO计算所必需的LLVM IR，然后调用LLVM从LLVM IR发出原生代码。

GPU后端目前通过LLVM NVPTX后端支持NVIDIA GPU，CPU后端支持多个CPU指令级架构（ ISA）。

4.3 支持平台

XLA目前支持x86-64和NVIDIA GPU上的JIT编译、以及适用于x86-64和ARM的AOT编译。

4.3.1 即时编译（JIT）

注：要使用JIT必须从源码编译包含XLA的TensorFlow。

为何使用即时编译（just-in-time compilation）
TensorFlow / XLA JIT编译器，可以通过XLA，编译和运行部分TensorFlow图（Graphs）。这比标准TensorFlow实现的好处是XLA可以将多个运算符（内核融合）融合到少量编译的融合的内核中。融合后运算符可以减少内存带宽并提高性能。

用XLA运行TensorFlow图
两种方法可以通过XLA运行TensorFlow计算：

将JIT编译的运算符放在CPU或GPU设备上；
将算子放在XLA_CPU或XLA_GPU TensorFlow设备上。

将算子直接放在TensorFlow XLA设备上，会强制算子在该设备上运行，主要用于测试。

注意：XLA CPU后端支持算子内的多线程并行（即它可以在多个计算核之间对单个算子进行分片），但它不支持不同算子间的并行（即不能跨多个计算核同时执行独立的算子）。XLA GPU后端与标准的TensorFlow实现性能不分高下，有时更快，有时更慢。

4.3.2 AOT编译

什么是tfcompile
tfcompile是一个独立的工具，提前（AOT）将TensorFlow图编译成可执行代码。它可以减少总二进制包的大小，还可以避免运行时开销。tfcompile的典型用例是将推理图编译成移动设备的可执行代码。

TensorFlow图通常由TensorFlow运行时执行，这会导致执行图中每个节点的运行时开销、更大的二进制包总大小，因为图本身的代码和TensorFlow运行时的代码，都需要可用。tfcompile生成的可执行代码不使用TensorFlow运行时，只依赖计算中实际使用的内核。

编译器构建在XLA框架之上。将TensorFlow桥接到XLA框架的代码，驻留在tensorflow / compiler下，其中还包括对TensorFlow图的即时（JIT）编译支持。

tfcompile有什么作用？
tfcompile接受一个子图，由TensorFlow的feed和fetches概念标识，并生成一个实现该子图的函数。feed是函数的输入参数，fetch是函数的输出参数。所有输入必须由feed完全指定；生成的修剪子图不能包含占位符或变量节点（ Placeholders and Variables Nodes）。通常将所有占位符和变量（ Placeholders and Variables）指定为feed，这可确保生成的子图不再包含这些节点。生成的函数打包为cc_library，其中包含导出函数签名的头文件，以及包含该实现的目标文件。用户编写代码以适当地调用生成的函数。

使用tfcompile
本节详细介绍了使用TensorFlow子图中的tfcompile生成可执行二进制文件的高级步骤。步骤是：

第1步：配置要编译的子图
第2步：使用tf_library构建宏来编译子图
第3步：编写代码来调用子图
第4步：创建最终的二进制文件

具体参考Using AOT compilation | XLA | TensorFlow。

5. Roadmap 2019

官方于2019年3月6日更新的2019年Lite开发规划，总的来说分为四部分，这里只提我认为比较关键的待开发特性：

易用性（Usability）
- 新转换器（New Convertor）。更好的处理graph转换（例如控制流、条件等等），替换掉TOCO；
- 端上训练支持；
- 更多operator支持；
性能（Performance）
- 更多端上设备支持；
- 支持Android NN API支持；
- 支持更多基于OpenGl和Metal的GPU kernel实现；
- 持续优化CPU在float和量化的模型；
优化（Optimization）
- 在后训练量化中，支持混合kernel和定点kernel；
- 更低bit-width的推理支持；
便携性（Portability）
- 单片机（Microcontroller）支持：增加8、16、32bit的微控制器架构在语音图像分类上的使用案例。

其中，除了持续优化已支持设备的kernel外，对于端上设备的训练、更低比特宽度的推理，都值得关注。

推荐阅读
嵌入式AI简报 (2020-01-06)

本作品采用知识共享署名-相同方式共享 4.0 通用许可协议进行许可。
欢迎关注公众号，关注模型压缩、低比特量化、移动端推理加速优化、部署。