TensorFlow Lite源码解析之一

1. 简介

TensorFlow Lite是一个用于端侧推理的工具集，它可以让我们将TensorFlow模型部署到手机、嵌入式设备甚至物联网设备上。它主要分为两部分：模型转换工具以及模型推理引擎。

顾名思义，模型转换工具是用于对模型进行转换的，目前也只是将TensorFlow导出的.pb模型转换为.tflite模型。对于高通SNPE这种面向多种框架导出的不同格式的模型的推理引擎，有必要将不同格式的模型统一到一个标准，因此需要一个转换工具，但是TensorFlow Lite是用于TensorFlow Lite模型推理的，那为什还要进行转换呢？原因很简单，TensorFlow Lite使用运行于移动端，有些设备甚至是资源非常有限，因此在内存和解析方面必须尽可能减少开销。TensorFlow导出的模型使用的是Protocol Buffer协议，因此有必要将它转换到性能更优的FlatBuffer格式。

作为端侧推理引擎，TensorFlow Lite有以下几个有点：

1. 低延迟，省去了与服务器沟通的来回时间；
2. 隐私性能好，所有数据都只在本地处理；
3. 随时随地，没网也行；
4. 低功耗，相比于网络通信，推理所需要的电能会更少。

接下来的日子，就让我们一起骑上心爱的小摩托，徜徉在TensorFlow Lite的代码里吧。

2. 使用流程

下面展示的是一个官方给的C++中使用TFLite的例子：

int main(int argc, char* argv[]) {
  if (argc != 2) {
    fprintf(stderr, "minimal <tflite model>\n");
    return 1;
  }
  const char* filename = argv[1];

  // Load model
  std::unique_ptr<tflite::FlatBufferModel> model =
      tflite::FlatBufferModel::BuildFromFile(filename);
  TFLITE_MINIMAL_CHECK(model != nullptr);

  // Build the interpreter with the InterpreterBuilder.
  // Note: all Interpreters should be built with the InterpreterBuilder,
  // which allocates memory for the Intrepter and does various set up
  // tasks so that the Interpreter can read the provided model.
  tflite::ops::builtin::BuiltinOpResolver resolver;
  tflite::InterpreterBuilder builder(*model, resolver);
  std::unique_ptr<tflite::Interpreter> interpreter;
  builder(&interpreter);
  TFLITE_MINIMAL_CHECK(interpreter != nullptr);

  // Allocate tensor buffers.
  TFLITE_MINIMAL_CHECK(interpreter->AllocateTensors() == kTfLiteOk);
  printf("=== Pre-invoke Interpreter State ===\n");
  tflite::PrintInterpreterState(interpreter.get());

  // Fill input buffers
  // TODO(user): Insert code to fill input tensors.
  // Note: The buffer of the input tensor with index `i` of type T can
  // be accessed with `T* input = interpreter->typed_input_tensor<T>(i);`

  // Run inference
  TFLITE_MINIMAL_CHECK(interpreter->Invoke() == kTfLiteOk);
  printf("\n\n=== Post-invoke Interpreter State ===\n");
  tflite::PrintInterpreterState(interpreter.get());

  // Read output buffers
  // TODO(user): Insert getting data out code.
  // Note: The buffer of the output tensor with index `i` of type T can
  // be accessed with `T* output = interpreter->typed_output_tensor<T>(i);`

  return 0;
}

从例子中可以看到，使用TFLite进行推理主要分为以下几步：

1. 加载模型文件。模型文件的内容使用类FlatBufferModel来管理，FlatBufferModel中有个一Model类型的成员，该成员由对应的Schema编译得到，它将根据FlatBuffer的解析规则将序列化的数据反序列化，关于FlatBuffer反序列化的原理在之前的文章《FlatBuffer内部解析原理简介》有介绍过，感兴趣的可以查看。Model拥有指针指向模型文件加载进入内存后所在的Buffer。这片Buffer如果是在支持mmap的操作系统中，便是一块与模型文件进行映射的内存；如果是在MCU等不支持mmap的嵌入式系统中，则直接是文件内容在内存中的一个拷贝。其关系如图所示；

2. 创建一个类Interpreter的实例。Interpreter是TFlite推理过程的核心，但是在实践中一般不直接使用其自身构造函数创建，而是使用建造者模式来创建其实例。当然了，一般情况下不使用，那么就是说不是不能用，那我们什么时候可以直接构造一个Interpreter的实例来使用呢？在使用场景非常简单的情况下，例如需要直接跑单个Op的时候。在这个过程中，推理引擎会对模型进行解析、分配模型运行过程中需要的内存空间以及确定模型中各个层之间的执行顺序；
3. 准备输入数据，这个需要根据模型要求的输入格式做对应的处理；
4. 调用Interpreter.Invoke()进行推理；
5. 获取推理结果。

总的来说，推理引擎工作过程基本都是一样的，分为模型加载-解析-推理这些过程，例如之间介绍过的ONNX Runtime也基本如此。

3. 实例详解

前面说过，作为TFLite推理过程中的核心，在实践中我们基本不会直接调用Interpreter这个类的构造函数来实例化，因为如果我们手动实例化Interpreter，那么我们就需要自己手动调用相关类去解析模型并填充到Interpreter对应的成员中。很显然，这么做是很繁琐并且容易出错的，并且要求使用者非常熟悉TFLite的内部结构，很显然这种要求是很不合理的，这就相当于你买了辆车，但是4S店给你的是一堆可以组成车的零件让你自己组装好再开。

正因如此，TFLite使用了建造者模式来应对这一繁琐的过程，调用者在正经推理的时候只需要创建一个InterpreterBuilder的实例，然后传递给它一个托管Interpreter的智能指针std::unique_ptr<Interpreter>，剩下的就交给时间，InterpreterBuilder会将蛋白质大分子水解成小分子的多肽和氨基酸，这一转变赋予了豆腐不一样的风味……额，不好意思，串台了。InterpreterBuilder会将加载进入到内存的模型反序列化得出其中所有的节点以及所包含的权重等信息，并为其分配对应的算子以及计算过程中所需的内存空间。

虽然实践中不直接实例化Interpreter，我们却可已通过直接使用Interpreter的一些例子来更好的理解TFLite的工作过程，正所谓管中窥豹，可见一斑。下面我们就来看一个例子。

template <class T>
void resize(T* out, uint8_t* in, int image_height, int image_width,
            int image_channels, int wanted_height, int wanted_width,
            int wanted_channels, Settings* s) {
  int number_of_pixels = image_height * image_width * image_channels;
  std::unique_ptr<Interpreter> interpreter(new Interpreter);

  int base_index = 0;

  // two inputs: input and new_sizes
  interpreter->AddTensors(2, &base_index);
  // one output
  interpreter->AddTensors(1, &base_index);
  // set input and output tensors
  interpreter->SetInputs({0, 1});
  interpreter->SetOutputs({2});

  // set parameters of tensors
  TfLiteQuantizationParams quant;
  interpreter->SetTensorParametersReadWrite(
      0, kTfLiteFloat32, "input",
      {1, image_height, image_width, image_channels}, quant);
  interpreter->SetTensorParametersReadWrite(1, kTfLiteInt32, "new_size", {2},
                                            quant);
  interpreter->SetTensorParametersReadWrite(
      2, kTfLiteFloat32, "output",
      {1, wanted_height, wanted_width, wanted_channels}, quant);

  ops::builtin::BuiltinOpResolver resolver;
  const TfLiteRegistration* resize_op =
      resolver.FindOp(BuiltinOperator_RESIZE_BILINEAR, 1);
  auto* params = reinterpret_cast<TfLiteResizeBilinearParams*>(
      malloc(sizeof(TfLiteResizeBilinearParams)));
  params->align_corners = false;
  params->half_pixel_centers = false;
  interpreter->AddNodeWithParameters({0, 1}, {2}, nullptr, 0, params, resize_op,
                                     nullptr);

  interpreter->AllocateTensors();

  // fill input image
  // in[] are integers, cannot do memcpy() directly
  auto input = interpreter->typed_tensor<float>(0);
  for (int i = 0; i < number_of_pixels; i++) {
    input[i] = in[i];
  }

  // fill new_sizes
  interpreter->typed_tensor<int>(1)[0] = wanted_height;
  interpreter->typed_tensor<int>(1)[1] = wanted_width;

  interpreter->Invoke();

  auto output = interpreter->typed_tensor<float>(2);
  auto output_number_of_pixels = wanted_height * wanted_width * wanted_channels;

  for (int i = 0; i < output_number_of_pixels; i++) {
    switch (s->input_type) {
      case kTfLiteFloat32:
        out[i] = (output[i] - s->input_mean) / s->input_std;
        break;
      case kTfLiteInt8:
        out[i] = static_cast<int8_t>(output[i] - 128);
        break;
      case kTfLiteUInt8:
        out[i] = static_cast<uint8_t>(output[i]);
        break;
      default:
        break;
    }
  }
}

上面的例子展示的是输入数据的尺寸不符合模型的输入尺寸要求，在数据处理过程中调用TFLite已实现的双线性插值算法对输入数据进行缩放的过程。

上面的代码虽然简单，但是它是一个完整的推理过程。当我们了通过这件简单的例子了解了整个推理过程以后，再去看其他的代码就会轻松许多。

首先，在代码开始的时候，手动创建了一个Interpreter实例并通过智能指针托管。整个创建过程知识往一个存储模型子图的列表里添加了一个空的子图；

紧接着，连续调用了两次AddTensor()这个方法，这个方法主要作用就在第一步创建的子图中预留对应数量了空间，用于后续张量数据的存放。这些张量用类TfLiteTensor表示，它们被保存在一个std::vecter列表里，并且结构体TfLiteContext中有一个指向其底层数组的指针。就好比要开会了，虽然与会者目前还没有来，但是准备会议的人需要知道这场会议有多少人参加，也好准备相应的物品：椅子少了，去搬点过来；多了就把多余的搬走；

接下来，通过SetInput(), SetOutput()这两个方法，告诉第一步创建的子图，第二部预留的张量中，哪些是这个子图的输入，哪些是这个子图的输出。这两个方法的参数都是一个包含表示张量位置的索引值的数组。就相当于用名牌指示哪些位置是哪些领导的；

然后，接着调用SetTensorParametersReadWrite()方法，初始化指定的张量，主要包括张量的类型（整型还是浮点型，多少位等）、名字、维度信息、是否可更改、初始值等方面。就相当于已经确定某某领导在某某位置，我们想在要开始针对不同的领导准备东西：喝什么茶、放不放麦等等；

继续往下，通过算子的名字和版本，查找对应的算子的实现。每一个算子都有一个TfLiteRegistration实例来表示，它是一个结构体，包含有函数指针指向对应算子的操作，如初始化、计算等。所有支持的算子都被保存在一个字典里，通过算子的名字以及版本号可以唯一确定。

再往下，需要指定一些超参数，因为在真正计算之前，所有这些参数都不可能传递给对应的函数，因此需要先存起来。在这里，TFLite通过一个叫做TfLiteNode的对象将这个算子的输入、输出、超参数等信息保存起来，并将它与对应算子捆绑在一起作为一个元素放在一列列表里。

到了这一步，我们好像已经把该准备的都准备好了，是不是可以开始做推理了?答案是不行，可能细心的读者已经发现，我们传递进来需要计算得数据还没有被保存到对应的TfLiteTensor里。但我们看到拷贝数据这一步很快就在接下来的代码中有体现。但是奇怪的是，在这之前还有一个AllocateTensor()调用，既然我们已经提前在一个数组里为张量分配了空间，为什么还有这个调用？
其实答案很简单：在列表中分配的空间，是用来存储张量的元数据的，而真正存储张量内容的空间至今还没有分配。正因如此，在数据拷贝之前，需要调用AllocateTensor()来分配所需的内存空间。

最终，调用Invoke()进行计算。当计算完成后，结果已经被保存在被指定为输出的张量中。

4. 总结

通过上面的解析，我们已经大概弄清楚了TFLite中的主要成员以及他们之间的关系。

1. Interpreter中可以包含多个SubGraph，每一个都代表原始模型中的一个子图；
2. 对于原始模型中的每一个节点都会有一个TfLiteNode来表示它，并且该节点会和一个TfLiteRegistration绑定在一起；
3. TfLiteNode负责保存输入、输出以及超参数等信息；TfLiteRegistration负责找到对应的算子的实现函数；
4. 所有的Tensor的元数据都保存在一个列表里，每个Tensor再通过指针获取实际存储数据的内存；

接下来的日子，我们将慢慢探索TensorFlow 的奥妙。

首发于个人微信公众号【爱码士1024】。微信扫描上方二维码或者微信搜索【爱码士1024】并关注，及时获取更多最新文章！
C++ | Python | Linux | 原理 | 源码，有一起玩耍的么？

Resources

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