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

阿里巴巴集团
原文链接: mp.weixin.qq.com

摘要

  本篇文章将从工程侧的角度出发,介绍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加密的方式采用秘钥异或的方式加密。秘钥可以设置为任意长度,但是为了安全起见建议设置较长的长度,例如 32位或者64 位。同时,考虑到加密长度如果过长,会导致客户端解密时间过久。因此,我们会设置最大的加密长度,例如1MB或者 512KB。

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组件等。欢迎持续关注~

扫码关注【闲鱼技术】公众号

 

文章分类
人工智能
文章标签