Tensorflow Lite自定义OP Pipeline的工程实践

摘要

本篇文章将从工程侧的角度出发，介绍Tensorflow Lite自定义OP Pipeline。本文将分为四个部分做具体介绍：概述、服务端自定义OP、客户端自定义OP和延展。首先，概述部分介绍什么是自定义OP。其次，服务端部分会介绍自定义OP如何注册、编译和加密。再次，客户端部分会介绍自定义OP如何下载验证、注册和运行。最后是延展部分，介绍现有自定义OP可以改进的方向。

图1自定义OP的特点

其中，本文的自定义OP具有如下四个特点：速度、能力、安全和创新。

I)速度：紧跟tensorflow Lite最新的版本，并修复了若干个问题，例如java nio bug，解决了文件落盘问题。

II)能力：能够动态下发模型和库，并且能够解决复杂问题。例如车牌识别，语音识别，获取封面帧，视频指纹等功能。

III)安全：自定义OP和模型都进行了加密处理，且解密的文件不会写入磁盘，从而保证了文件的安全性。

IV)创新：当前官方没有找到可以实现自定义OP客户端上使用的案例，本文创新性地采用了服务端编译三端的库文件，动态下发库文件供Android和IOS加载并使用的方式。具体的特点将在后续章节详细介绍。

关键词：

TensorflowLite,自定义OP,创新,工程实践

1.概述

1.1 Tensorflow Lite简介

本文基于tensorflow Lite是谷歌研发的人工智能学习系统，其命名来源于本身的运行原理。Tensor（张量）意味着N维数组，Flow（流）意味着基于数据流图的计算，TensorFlow为张量从流图的一端流动到另一端计算过程。TensorFlow是将复杂的数据结构传输至人工智能神经网中进行分析和处理过程的系统。

下图展示了tensorflow Lite的架构设计[1]：

图1.1 tensorflow lite架构图

1.2 OP简介

OP的输入/输出以Tensor的形式存在，可以由一个或者多个tensor构成[3]。使用内置的OP可以实现矩阵加法、乘法、池化、RELU、SOFTMAX等操作。在tensorflow Lite中现有内置的34种OP如下：

图1.2现有的内置OP类型

图1.2展示的均是内置OP，此外tensorflow Lite也支持自定义OP，允许用户添加特定的功能。自定义OP将会在第二章和第三章具体介绍。

2.服务端自定义OP

如图2.1所示，本章会在后续小节描述服务端自定义OP包括服务端自定义OP代码注册、生成自定义OP节点加入并生成Lite模型、OP加密等实现方法。

图2.1服务端自定义OP流程图

2.1自定义OP注册

本小节介绍服务端自定义OP代码如何注册。以谷歌官方ZeroOut为例，需要引入tensorflow相关的头文件。REGISTER_OP这个宏来定义自定义OP的名字，Input为输入格式(int32)，Output为输出格式(int32)。

图2.2自定义OP服务端注册

之后，需要在compute函数中实现对应的OP节点运行时的逻辑。也需要在注册文件的最后声明使用CPU或者GPU编程。如图2.3和图2.4所示。

图2.3自定义OP服务端实现函数

图2.4服务端自定义OP定义使用方式（GPU or CPU）

2.2服务端自定义OP编译

服务端会实现三端（服务端、Android和IOS）的库文件编译，Android编译采用了ndk-build编译，IOS和服务端采用Makefile编译。

自定义OP的编译前提是需要配置tensorflow环境[4]，这不在本文中详细展开了。配置完成后，引入相关的tensorflow头文件，服务端的OP编译执行如下命令即可[5]。

图2.5服务端自定义OP编译方式

考虑到ndk-build在Android端有较好地优化[6]，因此Android端编译使用ndk-build的方式编译。具体方式如图2.6所示。考虑到armeabi-v7a会引入ARM指令集优化[2]，因此优先采用该指令。使用--gc-sections -ffunction-sections-fdata-sections命令不引入没有使用的函数和数据，可以减小库文件体积。

图2.6 Android端自定义OP编译方式

在IOS端，可以使用Makefile编译，但是需要使用lipo工具将不同平台的库(例如armv7

，armv7s和arm64等)打包集成在一起。具体用法详见图2.7。

图2.7IOS端自定义OP编译方式

2.3生成自定义OP模型

前一小节介绍了服务端如何注册并编译生成自定义OP。这一节会介绍如何将自定义OP节点加入Lite模型中。

当生成zero_out.so时，可以利用tf.load_op_library函数将zero_out函数引入到模型中来，如图2.8所示：

图2.8生成自定义OP模型

2.4自定义OP加密

为了保证自定义OP从服务端下发的安全性，需要将自定义OP加密。本文中自定义OP加密的方式采用秘钥异或的方式加密，具体方式如图2.9所示。秘钥可以设置为任意长度，但是为了安全起见建议设置较长的长度，例如32位或者64位。同时，考虑到加密长度如果过长，会导致客户端解密时间过久。因此，我们会设置最大的加密长度，例如1MB或者512KB。

图2.9 自定义OP加密方式

3.客户端自定义OP

服务端生成自定义OP后，客户端需要下载相应的模型和自定义OP库文件，之后注册该自定义OP并加载模型即可在客户端使用自定义OP实现的功能。

图3.1客户端自定义OP流程图

3.1模型下载校验

由于模型和库文件下载过程中可能会产生中断、失败或者版本错误问题，因此客户端引入MD5校验来验证模型下载的正确性：客户端和服务端下发文件的MD5值相同，则认为下载正确，选择加载该文件，将自定义OP注册到tensorflowLite。

3.2自定义OP注册

自定义OP注册使用服务端相同的秘钥解密。为了安全起见，解密只保存在内存中（android采用了nio方式读取），也就是文件不落盘。然后使用如下函数注册：

AddCustom(constchar* name, TfLiteRegistration* registration)

其中，name表示输入自定义OP的名字，比如ZeroOut，GetCover等，registration表示需要注册TfLiteRegistration的函数，具体函数如下：

I)void*(*init) (TfLiteContext* context, const char* buffer, size_t length)

II)void(*free) (TfLiteContext* context, void* buffer)

III)TfLiteStatus(*prepare) (TfLiteContext* context, TfLiteNode* node)

IV)TfLiteStatus(*invoke) (TfLiteContext* context, TfLiteNode* node）

图3.2自定义OP注册函数

以上四个函数，分别是init，free，prepare和invoke函数，一般实现prepare预处理函数逻辑和invoke执行函数逻辑即可。

在IOS可以直接调用AddCustom函数即可，在Android端并没有提供这个接口，需要在JNI层暴露这个函数，具体如图3.3所示：

图3.3 AddCustom动态加载OP函数

此外，该函数还采用了dlopen的方式实现动态加载。例如当服务端动态下发了名为ZeroOut的OP时，只需要知道ZeroOut.so的路径，即可调用Prepare和Invoke方法实现动态注册OP，不需要安装新的APK。

当自定义OP注册完成后，即可通过tensorflowLite的模型加载运行自定义OP的功能[7]。

4.延展

本文的工程实践包括自定义OP框架在服务端上编译，并在客户端使用自定义OP的功能。自定义OP可以很好地补充tensorflowLite实现功能较单一的缺陷。然而，本文的自定义OP也存在很多不足，可以有如下优化：

1）没有压缩相关的库文件和模型文件，下发文件可能较大。因此，可以引入模型压缩算法。

2）现有的格式只有int32、float32等几种单一的格式，可以引入string，double64等数据格式增加可用性。

参考文献：

[1]https://www.tensorflow.org/mobile/tflite

[2] https://baike.baidu.com/item/ARM指令集/907786?fr=aladdin

[3] https://www.tensorflow.org/api_docs/python/tf/Session

[4] https://www.tensorflow.org/versions/r1.5/install/install_sources

[5] https://www.tensorflow.org/extend/adding_an_op

[6] https://developer.android.google.cn/ndk/guides/ndk-build

[7] https://www.tensorflow.org/versions/r1.5/mobile/tflite/demo_android

后续我们将会推出一系列基于tensorflow的工程实践项目。客户端会推出各类宝贝识别，宠物识别等。服务端会推出UI自动化项目，自动识别各种业务组件，各类UI组件等。欢迎持续关注～