TensorFlow Lite源码解析之三（委托）

1. 前言

所谓劳心者治人，劳力者治于人（所以“劳力士”又叫“打工人”？）。小农经济下自己什么都干，虽然饿不死，但是生产效率也肯定非常低下；只有让专业的人做专业的事，才能最大限度发挥每个人的价值。同理，对于一个推理引擎而言，完成一次计算，可以有多种选择，从头到尾都让CPU做运算理论上也不是什么问题。但是除了结果正确之外，还有另外一项关键指标：时间！密集计算确实非CPU所长，因此他该做的是协调资源，脏活累活让GPU、DSP这些替他干。一个精诚合作的团队力量是非常强大的。那么，TensorFlow Lite是怎么利用设备上的硬件加速器做加速的呢？今天我们就来掰扯掰扯。

2. 正文

对于加速器的使用，不同的推理引擎实现的方法不一样，叫法也不统一：例如ONNX Runtime中称之为 Provider，而在TensorFlow Lite中则称之为 Delegate。当然，叫法其实不重要，下厨和烧菜本身都是指做东西吃，我们更关心的是制作的过程，万一哪天没外卖吃了要自己动手呢？

作为对比，我们可以简单回忆一下ONNX Runtime的做法（感兴趣的可以直接戳这里）：ONNX Runtime默认的情况下会自动识别机器上可用的加速器，每个已识别的加速器都会声明其自身支持的算子，系统根根这些声明进行网络分片。对于多个加速器都声明支持的算子，则按照预设的加速器优先级来切分网络。

对于TensorFlow Lite，大致原理也差不离，只不过某些方面稍有不同。

首先，TensorFlow Lite默认情况下并不会去识别机器上所拥有的加速器，用户必须在代码中显式使用。下面我们将会以Hexagon Delegate为例，讲解Delegate的使用。

在TFLite中要想使用加速器，用户需要在加载网络模型、使用建造者模式获得Interpreter的实例以后，必须调用与Hexagon 相关的API创建一个 Hexagon Delegate，之后，调用interpreter->ModifyGraphWithDelegate(delegate)注册Hexagon Delegate。只有这样，TFLite在推理的过程中才会使用Hexagon极速器进行加速。示例代码如下所示。

#include "tensorflow/lite/delegates/hexagon/hexagon_delegate.h"

// Assuming shared libraries are under "/data/local/tmp/"
// If files are packaged with native lib in android App then it
// will typically be equivalent to the path provided by
// "getContext().getApplicationInfo().nativeLibraryDir"
const char[] library_directory_path = "/data/local/tmp/";
TfLiteHexagonInitWithPath(library_directory_path);  // Needed once at startup.
::tflite::TfLiteHexagonDelegateOptions params = {0};
// 'delegate_ptr' Need to outlive the interpreter. For example,
// If use case will need to resize input or anything that can trigger
// re-applying delegates then 'delegate_ptr' need to outlive the interpreter.
auto* delegate_ptr = ::tflite::TfLiteHexagonDelegateCreate(¶ms);
Interpreter::TfLiteDelegatePtr delegate(delegate_ptr,
  [](TfLiteDelegate* delegate) {
    ::tflite::TfLiteHexagonDelegateDelete(delegate);
  });
interpreter->ModifyGraphWithDelegate(delegate.get());
// After usage of delegate.
TfLiteHexagonTearDown();  // Needed once at end of app/DSP usage.

上面的代码大致可分成两个阶段：

1. 创建Delegate，对Hexagon 进行一些初始化工作；
2. 模型修改，实际上就是模型切片，这一步确定了模型中哪些算子会使用Hexagon进行加速。

接下来会详细讲解这两个过程。

2.1. Delegate 的创建

Hexagon Delegate的创建过程，也可以分成两部分。

1. 首先，通过指定路径寻找并加载相关的动态链接库，寻找的路径包括用户指定的路径以及/system/lib/rfsa/adsp;/system/vendor/lib/rfsa/adsp;/dsp这三个高通指定的路径。然后获取到所有构建、执行、销毁Hexagon NN网络这一整个过程所需的所有API的句柄，将这些句柄封装在一个HexagonNN的结构体中。最后调用一个名叫hexagon_nn_global_init()的函数去初始化与Hexagon NN相关的部件，至于具体怎么初始化的，就不得而知了，这属于高通的私有代码。实际上，我们也不需要知道。
2. 第二步，创建一个TfLiteDelegate结构体，这个结构体中主要包含了加速器环境准备、加速器与CPU之间的数据交换等操作的函数指针。但是在这一步中最重要的是，为.data_这个成员赋值：data_成员会指向一个类名为HexagonDelegate的实例，其主要提供的是判断哪些算子是可以在Hexagon上面进行计算以及获取一个表示Hexagon算子的接口。他们之间的关系如下图所示。在后续，TFLite便会通过这一步中Prepare指针所指向的函数，其实也就是位于simple_delegate.cc的DelegatePrepare()函数。

Delegate组件的关系

到此，Delegate的创建就算是完成了。总结下来，其实也没做什么实质性的工作，主要就是让各个API之间各就各位，为下一步做准备。

2.2. 模型切片

对模型的切片是在调用interpreter->ModifyGraphWithDelegate(delegate)的时候完成了。如下图所示，图中展示了ModifyGraphWithDelegate(delegate)的调用链关系，并且使用不同颜色对不同的函数归属做出了区分。

image.png

从上图中可以看到，对于模型分片，主要也是分成三步走：

1. 首先就是确定模型中哪些节点是可以在该加速器上加速的。Partition()会调用到指定的的Delegate中的IsNodeSupportedByDelegate()方法，这个方法不仅会判断节点的类型，也会判断节点的参数是否符合该加速器的计算要求。例如目前Hexagon Delegate要求很多节点的输入为位宽为8位的量化数据，如果模型未经量化，不好意思，出门左转慢走不送。Partition()最终会返回一个数组，这个数组的元素为Hexagon 所支持的节点的id。这个数组会最终最为参数传递到ReplaceNodeSubnetsWithDelegateKernels()中，用于指示将哪些节点原有绑定的算子给替换成跑在Hexagon上的算子；
2. 第二步，调用GetDelegateKernelRegistration()获取Hexagon的算子，通过函数指针来实现；

3. 第三步，做算子替换，并且会从新做内存分配调整。替换的方法是按照模型的拓扑结构从头开始，将模型分成多个子网络，每个子网络中包含的都是当前Delegate所支持的算子，并且它们在原模型的拓扑结构上是连续的。例如下图，假设Hexagon只支持卷积以及PReLU操作，那么整个模型就会被分成四个子网络（两个分支无论谁先谁后都会在Concat之前执行完，所以在拓扑上他们是连续的）。两个红框内的节点会在Hexagon上执行，因此红框内的字网络会替换成一个表示为Hexagon Kernel的节点，最终从CPU的角度看，整个模型从12个节点变成了4个节点，如下图所示。

   
   
   替换前模型示意图

替换后的模型示意图

在模型替换后，TFLite也会对内存重新分配，在CPU这边只会留下Hexagon子网络的输入输出张量，其他子网络内部的张量的分配，就由Hexagon自行决定，这也是高通的私有代码，我们无从知晓了。

3. 疑问

至此，TFLite如何获取加速器做加速的流程我们已经大致弄清楚了，但是还有一些疑问：

1. 我可以委托多个加速器做加速么？
2. 对于多个加速器都支持的算子，怎么确认使用哪一个？

第一个问题，很显然是可以通过多个Delegate使用多个加速器做加速的，只要我们在Interpreter实例化之后以及调用Interpreter.Invoke()之前依次初始化不同类型的Delegate，例如GPU、NNAPI等的Delegate，并且使用interpreter->ModifyGraphWithDelegate(delegate)对这些Delegate一一注册便可以。值得一提的是，TFLite会自动注册一个名叫FlexDelegate的Delegate，只不过这个Delegate也是在CPU上执行而已。

既然可以使用多个Delegate，那么肯定会出现多个算子同时被多个Delegate支持的情况，那么怎么解决？其实通过上面的介绍，我们已经清楚了，也是与ONNX Runtime类似——先到先得。由于每个Delegate都会把原有模型的网络中支持的子网络用一个表示自身类型的节点代替，因此之后的Delegate是不会再看到已被之前的Delegate替换了的节点的，例如有三个连续的卷积操作节点类型为Conv，他们被Hexagon Delegate替换成了一个类型为kHexagonKernel节点，后续的Delegate再去查看整个网络结构，只能看到一个类型为kHexagonKernel的节点，是不会看到那三个 Conv节点的。

4. 总结

通过上面的分析我们知道TFLite调用加速器的大概流程如下：

1. 生成一个Delegate实例，所有的Delegate都继承自SimpleDelegateInterface，其作用主要有两点：判断模型中那些节点可以在对应加速器上运算以及获取操纵该加速器的函数的句柄；
2. 划分模型子网络，对于可运行的在某个加速器上的子网络，使用一个标示性的节点替换掉其中所有节点，替换掉的节点如何计算以及内存如何分配又对应加速器决定，CPU只负责其输入输出张量的分配以及对应数据的传输；
3. 为第二步中标示性节点绑定操作对应的加速器的句柄，这样CPU便在模型执行到该节点时便可通知对应加速器开始计算。

接头暗号：爱码士1024

欢1迎2关3注4个5人6微7信8公9众10号：爱码士1024

5. Resources

[1] https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite