APP下载

TensorFlow Lite:端侧机器学习框架

2020-09-24李双峰

计算机研究与发展 2020年9期

李双峰

(Google TensorFlow团队 北京 100190)shuangfeng@google.com)

TensorFlow Lite[1](TFLite)是一个轻量、快速、跨平台的专门针对移动和IoT应用场景的机器学习框架,是开源机器学习平台TensorFlow[2-4](TF)的重要组成部分.它致力于“一次转换,随处部署”,支持安卓、iOS、嵌入式Linux以及MCU等多种平台,降低开发者使用门槛,加速端侧机器学习(on-device machine learning, ODML)的发展,推动机器学习无处不在.本文介绍了工业界端侧机器学习的最新趋势以及TFLite如何加速其发展,包括:

1) 端侧机器学习的趋势、挑战和典型应用,以及TFLite的起源;

2) TFLite的系统架构;

3) TFLite的最佳实践以及适合初学者的工具链;

4) 未来的发展方向.

1 TensorFlow Lite推动端侧机器学习的发展

1.1 机器学习的普及以及TensorFlow的发展

机器学习的发展改变着语音交流、机器翻译、健康医疗和城市交通等我们生活中的诸多领域,越来越多的企业、组织和个人尝试用机器学习解决业务或者生活中遇到的难题.开发者数量的大量增加也意味着需要降低技术门槛,让更多人参与进来.

在2015年底,Google开源了端到端的机器学习开源框架TensorFlow[2-4]:它既用于研究,也用于大规模生产领域;既支持大规模的模型训练,也支持各种环境的部署,包括服务器和移动端的部署;支持多种语言比如Python,C++,Java,Swift甚至JavaScript.TensorFlow提供了全面灵活的工具生态,帮助解决各种挑战性的问题.而近年来移动化浪潮和交互方式的改变,使得机器学习技术开发也在朝着轻量化的端侧发展,于是TensorFlow团队在2017年底开源了TFLite[5],一个轻量、快速、兼容度高的专门针对移动应用场景的深度学习工具,降低了端侧深度学习技术的门槛.

1.2 端侧机器学习的机遇和挑战

伴随移动和IoT设备的普及,世界以超乎想象的方式连接在一起.如今已有超过30亿的智能手机用户[6],以及超过70亿的联网IoT设备[7](不包括手机、电脑等).手机成本不断降低,并且随着微控制器(mircrocontrollers, MCU)和微机电系统(micro-electro-mechanical systems, MEMS)的发展,高性能低功耗的芯片使得“万物”智能具有了可能性.从智能穿戴、智能家居到共享单车,从工业控制到车载设备,这些设备都有了智能化的基础.我们有时把这些设备统称为边缘设备(edge device).

将边缘设备的数据传输到云端处理很多时候不是最经济有效的方式,它带来了延迟,降低了复杂网络条件下的可靠性,也引起了更多的隐私顾虑,从而影响用户体验.用户对交互的需求越来越高,快速、及时的智能反应是消费者的普遍期待,以智能音箱为例,唤醒的响应速度是良好体验的基础.

我们把在边缘设备上的运行机器学习统称为端侧机器学习(ODML),它为万物智能互联带来了新的机遇:

1) 更快、更紧密的交互方式,因为模型在本地执行的延迟小.比如及时的语音唤醒、直播视频的实时图像分割或者辅助驾驶的目标检测,都需要在本地执行.

2) 在复杂网络环境下仍可提供可靠服务,比如在网络基础设施相对落后的国家或偏远地区,在复杂环境下(比如隧道中),带宽可能有限或无法联网.

3) 更好的保护隐私,因为在本地进行数据收集和处理,减少了数据上传.

然而实现端侧机器学习有很多挑战,因为边缘设备:

1) 算力有限,限制了模型的复杂度.

2) 内存有限,限制了模型的大小.

3) 电池有限,需要模型运算效率更好.比如对智能手表而言,省电非常关键.

通常,移动设备上耗电最多的是无线电,如果在本地执行,尤其是有DSP或NPU等硬件加速器的情况下,电池续航时间就更长.

4) 计算硬件生态碎片化严重,比如CPU,GPU,DSP,NPU等,如何将这些异构硬件真正利用起来是一大难题.这与云端不同,在云端,硬件加速器供应商通常非常集中,如NVIDIA GPU或TPU.

1.3 TensorFlow Lite的起源和相关工作

TFLite是在边缘设备上运行TensorFlow模型推理的官方框架,它跨平台运行,包括Android,iOS以及基于Linux的IoT设备和微控制器.

TensorFlow适用于云端的大型、大功率设备,以及本地的工作站设备.当边缘设备的需求增加时,我们尝试简化TensorFlow并在移动设备上运行,这就是TF Mobile项目:它是一个缩减版的TensorFlow变体,简化了算子集,也缩小了运行库(runtime).然而,我们始终难以大大缩小运行库;同时运行库的扩展性方面也存在问题,如何将其映射到移动环境中所用的各种异构加速器上困难重重.

TFMini是Google内部用于计算机视觉场景的解决方案,它提供了一些工具(比如TOCO转换工具),压缩模型(比如删掉训练相关的不必要节点),进行算子融合并生成代码.它将模型嵌入到二进制文件中,这样我们就可以在设备上运行和部署模型.TFMini针对移动设备做了很多优化,性能优秀,但在把模型嵌入到实际的二进制文件中时兼容性存在较大挑战,因此TFMini并没有成为通用的解决方案.

基于TF Mobile的经验,也继承了TFMini和内部其他类似项目的很多优秀工作,我们设计了TFLite:

1) 更轻量.在32 b安卓平台下,核心运行时的库大小只有100 KB左右,加上支持基本的视觉模型(比如InceptionV3和MobileNet)所需算子时,总共300 KB左右,而使用全套算子库时,只有1 MB左右.

2) 特别为各种端侧设备优化的算子库.

3) 能够利用各种硬件加速,比如GPU,DSP等.

TFLite兼具性能和通用性,已经完全取代了TF Mobile和TFMini,成为TensorFlow针对移动和IoT设备的官方框架,被广泛使用.

其他相关工作包括:CoreML[8]是iOS平台的解决方案,而TFLite强调其优秀的跨平台能力.中国的开源端侧机器学习框架代表包括腾讯NCNN[9]、小米MACE[10]和阿里巴巴MNN[11]等,它们专注于移动推理平台,而TFLite是TensorFlow生态的一部分,支持从训练到多种平台部署[3](比如TFX,TF.js,TFLite),并提供完整的工具链(比如TensorBoard).

1.4 端侧机器学习应用的蓬勃发展

端侧机器学习在图像、文本和语音等方面都有非常广泛的想象空间.全球有超过40亿[12]的设备上部署着TFLite,这个数字还在不断增加当中.Google的大量产品部署着TFLite,比如Google Assistant,Google Photos等;国际巨头比如Uber,Airbnb等[12],以及国内的许多大公司,比如网易、爱奇艺和WPS等,都在使用TFLite.

一方面,端侧机器学习能解决的问题越来越多元化,这带来了应用的大量繁荣.前期在图像和视频方面广泛应用,比如Google Lens[13],Google Photos,Google Arts & Culture[14].最近,离线语音识别方面有很多突破,比如Google Assistant宣布了完全基于神经网络的移动端语音识别[15],效果和服务器端十分接近,服务器模型需要2 GB大小,而手机端只需要80 MB.端侧语音识别非常有挑战,它的进展代表着端侧机器学习时代的逐步到来.一方面依赖于算法的提高,另外一方面TFLite框架的高性能和模型优化工具也起到了很重要的作用.离线语音可以带来很好的用户体验,比如无需网络环境即可使用,无隐私顾虑.Google Pixel 4手机上发布了Live Caption[16],自动把视频和对话中的语言转化为文字,大大提高了有听力障碍人群的体验(accessibility).更进一步,端侧自然语言处理将会有巨大的前景.我们发布了基于MobileBERT模型的问题回答系统[17-18],速度非常快,在普通CPU上可实时响应,创造了无缝的用户体验.

另外一方面,模型越来越小,无处不再.Google Assistant的语音功能部署在非常多元的设备上,比如手机端、手表、车载和智能音箱上,全球超过10亿设备.更激动人心的前景发生在IoT领域,TFLite可以支持微控制器(microcontrollers, MCU)[19],而MCU是单一芯片的小型计算机,没有操作系统,只有内存,也许内存只有几十KB.很多设备上都有MCU,全球有超过一千五百亿的MCU.MCU低功耗、便宜、无处不在.TFLite发布了若干MCU上可运行的模型,比如识别若干关键词的语音识别模型和简单的姿态检测模型,模型大小都只有20 KB左右,我们基于此可构建很有意思的应用,比如更智能的小玩具:它在MCU上运行上述模型,呼叫设定的特定呢称时就会发出某个声音,而当拿起玩具做一些设定的动作时就会响应.当随处可见的物品都在MCU上部署机器学习的模型时,智能开始无处不在.

在中国,在移动应用方面,网易使用TFLite做OCR处理[20];爱奇艺使用TFLite来进行视频中的AR效果[21],而WPS用它来做一系列文字处理[22].在IoT方面,出门问问智能音箱使用TFLite来做热词唤醒[23](对于智能音箱而言,准确、实时、轻量化低功耗的唤醒非常关键),科沃斯扫地机器人使用TFLite在室内避开障碍物[24].另外,TFLite也非常适合工业物联智能设备的开发,因为它很好地支持如树莓派及其他基于Linux SoC的工业自动化系统.创新奇智应用TFLite开发智能质检一体机、智能读码机等产品,应用到服装厂质检等场景[25].

2 TensorFlow Lite系统架构

我们设计TFLite,目标是:

1) 轻量.缩小运行库和模型大小,减少内存消耗,适用于更多设备.

2) 高性能.针对移动和IoT设备深度优化,利用多种硬件加速机器学习,比如利用ARM CPU最新指令、GPU、DSP和NPU.

3) 跨平台、兼容度高.支持安卓、iOS、嵌入式Linux以及MCU等多种平台,支持多种“一次转换,随处部署”.

4) 易用.与TensorFlow紧密集成,实现从训练到部署过程流畅,提供丰富的平台相关API,提供丰富的模型库、完整的实例和文档,以及丰富的工具链,降低开发者使用门槛.

5) 可扩展性.模块化,易定制,容易扩展到更多硬件支持,定制更多算子.

Fig.1 Main components of TFLite图1 TFLite的主要组成部分

图1展示了TFLite的主要组成部分:

1) TFLite模型转换器[26](converter).TFLite自带一个转换器,它可以把TensorFlow计算图,比如SavedModel或GraphDef格式的TensorFlow模型,转换成TFLite专用的模型文件格式,在此过程中会进行算子融合和模型优化,以压缩模型,提高性能.

2) TFLite解释执行器(interpreter).进行模型推理的解释执行器,它可以在多种硬件平台上运行优化后的TFLite模型,同时提供了多语言的API,方便使用.

3) 算子库(op kernels).TFLite算子库目前有130个左右,它与TensorFlow的核心算子库略有不同,并做了移动设备相关的优化.

4) 硬件加速代理(hardware accelerator delegate).我们将TFLite硬件加速接口称delegate(代理),它可以把模型的部分或全部委托给另一个硬件后台执行,比如GPU和NPU.

图2展示了在TensorFlow 2.0中TFLite模型转换过程,用户在自己的工作台中使用TensorFlow API构造TensorFlow模型,然后使用TFLite模型转换器转换成TFLite文件格式(FlatBuffers格式).在设备端,TFLite解释器接受TFLite模型,调用不同的硬件加速器比如GPU进行执行.

Fig. 2 TFLite model conversion图2 TFLite模型转换过程

2.1 TensorFlow Lite模型转换器

图2中描述了模型转换过程,转换器可以接受不同形式的模型,包括Keras Model和SavedModel(TF 2.0中推荐的格式),开发者可以用tf.Keras或者低层级的TensorFlow API来构造TensorFlow模型,然后使用Python API或者命令行的方式调用转换器(推荐使用Python API,更灵活).比如:

1) Python API.调用tf.lite.TFLiteConverter,可用TFLiteConverter.from_saved_model(),或TFLiteConverter.from_keras_model();

2) 命令行:tflite_convert--saved_model_dir=tmpmobilenet_saved_model--output_file=tmpmobilenet.tflite

在TF 1.x版本中,还支持GraphDef格式,如果需要,请使用:tf.compat.v1.lite.TFLiteConverter.

转换器做了2类优化工作:

1) 算子优化和常见的编译优化,比如算子融合、常数折叠(constant folding)或无用代码删除等.TFLite实现了一组优化的算子内核,转化成这些算子能在移动设备上实现性能大幅度提升.比如让我们将Relu融合到卷积等高级算子中,或优化LSTM算子.

2) 量化的原生支持.在模型转换过程中使用训练后量化(post-training quantization)非常简单,不需要改变模型,最少情况只需多加一行代码,设置converter.optimizations=[tf.lite.Optimize.DEFAULT].

2.2 TensorFlow Lite FlatBuffers格式

TFLite模型文件格式采用FlatBuffers[27].与Protocol Buffers[28]类似,FlatBuffers是Google的一个开源的跨平台序列化格式,最初为视频游戏而设计.它在开发过程中更注重考虑实时性,内存高效,这在内存有限的移动环境中是极为关键的.它支持将文件映射到内存中(mmap),然后直接进行读取和解释,不需要额外解析.我们将其映射到干净的内存页上,减少了内存碎片化.另外,相对于Protocol Buffers,它有更小的依赖,因此也会减小二进制文件大小.

TFLite代码中schema.fbs[29]文件使用FlatBuffers定义了TFLite模型文件格式,我们摘取了其中一些关键样例代码,以帮助理解模型所包含的信息,如图3所示.TFLite模型文件是一个层次的结构:

1) TFLite模型(model)由子图(subgraph)构成,同时包括用到的算子库和共享的内存缓冲区.

2) 张量(Tensor,多维数组)用于存储模型权重,或者计算节点的输入和输出,它引用Model的内存缓冲区的一片区域,提高内存效率.

3) 每个算子实现有一个OperatorCode,它可以是内置的算子,也可以是自定制算子,有一个名字.

4) 每个模型的计算节点(operator)包含用到的算子索引,以及输入输出用到的Tensor索引.

5) 每个子图包含一系列的计算节点、多个张量,以及子图本身的输入和输出.

Fig. 3 TFLite schema.fbs samples[29]图3 TFLite schema.fbs样例代码[29]

模型转换器最早来源于TFMini项目,称为TOCO转换器.最近我们基于Google最新的机器学习编译技术MLIR[30]重写了转换器,将算子的TensorFlow dialect映射到算子的TFLite dialect.它提供了更好的转换错误追踪和调试功能,也支持了更多的新模型(比如Mask R-CNN,Mobile BERT),特别是对控制流(control flow)有更好的支持.另外,也让TFLite转换器具有更好的可扩展性,同时可以利用机器学习编译技术的最新成果.

2.3 TensorFlow Lite解释执行器

TFLite解释执行器[31]针对移动设备从头开始构建,具有3个特点:

1) 轻量级.在32 b安卓平台下,编译核心运行时得到的库大小只有100 KB左右,如果加上所有TFLite的标准算子,编译后得到的库大小是1 MB左右.它依赖的组件较少,力求实现不依赖任何其他组件.

2) 快速启动.既能够将模型直接映射到内存中,同时又有一个静态执行计划,在转换过程中,我们基本上可以提前直接映射出将要执行的节点序列.采取了简单的调度方式,算子之间没有并行执行,而算子内部可以多线程执行以提高效率.

3) 内存高效.在内存规划方面,采取了静态内存分配.当运行模型时,每个算子会执行prepare函数,它们做必要的内存分配.我们会分配一个单一的内存块,而这些张量会被整合到这个大的连续内存块中.不同张量之间甚至可以复用内存以减少内存分配.

参照图3中TFLite模型格式,我们会执行模型的子图,而根据数据依赖关系,子图中的计算节点已经被提前静态规划好,保证每个计算节点在执行前它所需要的张量已经被之前的计算准备好.所以,总体上依次执行即可.

使用解释执行器通常需要包含4步:

1) 加载模型.将TFLite模型加载到内存中,该内存包含模型的执行图.

2) 转换数据.模型的原始输入数据通常与所期望的输入数据格式不匹配.例如,可能需要调整图像大小,或更改图像格式,以兼容模型.

3) 运行模型推理.使用TFLite API执行模型推理.

4) 解释输出.解释输出模型推理结果.比如,模型可能只返回概率列表,而我们需要将概率映射到相关类别,并将其呈现给最终用户.

TFLite提供了多种语言的API,正式支持的有Java,C++和Python,实验性的包括C,Object C,C#和Swift.大家可以从头自己编译TFLite,也可以利用已编译好的库,Android开发者可以使用JCenter Bintray的TFLite AAR,而iOS开发者可通过CocoaPods在iOS系统上获取.

Fig. 4 TFLite Java API code samples图4 TFLite Java API代码样例

图4是使用Java API的一个示例:创建输入缓冲区和输出缓冲区,填充输入,并定制推理选项(比如是否使用GPU).之后创建TFLite执行器,输入TFLite模型并执行.

2.4 高性能算子库和可扩展性

目前TFLite约有130个算子(op kernels),大多同时支持浮点和量化类型,可以在这里找到所有支持的算子[32].很多算子专门针对ARM CPU进行了优化,包括基于汇编和intrinsics级别的NEON指令集的多种优化.针对矩阵相乘,算子之前浮点运算使用Eigen,量化运算使用gemmlowp[33],最近我们重新设计了一个叫Ruy[34]的高性能矩阵相乘库,统一了浮点运算和量化运算,并且性能有较大提升.特别针对一些重度使用的算子做了深度优化,比如涉及到卷积或者LSTM的算子.另外,不少算子支持多线程执行.

算子的定义基于C语言接口TfLiteRegistration,包括4个函数(init,free,prepare,invoke),开发者可方便自定义算子.

在创建解释器时,用户可以提供自定义的算子解析器(OpResolver),从而控制模型执行所用到的算子实现,如图5所示.使用内置算子解析器(BuiltinOpResolver)时使用默认的算子.每个算子都有自己的版本控制.我们还提供机制,把不用的算子自动删掉.

Fig. 5 TFLite customized op图5 TFLite自定义算子

当遇到不能支持的算子时,可以自定义算子,另外一种选择是,复用TensorFlow算子,称为Select TF Ops,只需要多加一行转换器参数就可以开启:

converter.target_spec.supported_ops=[tf.lite.OpsSet.TFLITE_BUILTINS,tf.lite.OpsSet.SELECT_TF_OPS]

它可能带来的问题是所依赖的运行库更大,同时TensorFlow算子在移动设备中未必足够优化.

2.5 CPU性能

CPU最具普适性,其性能最为关键,因此我们持续不断提升CPU性能.矩阵运算(GEMM[35])效率对于深度学习非常关键,过去我们使用gemmlowp[33]库进行量化矩阵乘法,使用Eigen库进行浮点乘法.最近,我们从头开发了一个针对移动环境CPU特别优化的矩阵乘法库Ruy[34],统一了浮点运算和量化运算.之前面向桌面和云端的矩阵乘法库则更加专注于大矩阵乘法的峰值性能,而Ruy对各种矩阵大小的乘法都表现优异:不仅适合大矩阵,也适合当前TFLite应用中最常见的矩阵运算(往往是小矩阵或者多种多样的矩阵形状).Ruy目前支持浮点数和8 b整型,主要针对ARM CPU(64 b和32 b),而我们也在针对英特尔x86架构进行优化.

自TensorFlow 1.15版本开始,Ruy在所有ARM设备上默认启用,大范围地加速了模型,特别是那些卷积比较多的视觉模型,将延迟降低了1.2~5倍,尤其在具有NEON点积内置函数的硬件上.图6以最常见的MobileNet V1为例,浮点模型和量化模型都有明显提高.

Fig. 6 Single thread CPU on Pixel4 图6 Pixel 4上单线程CPU性能

性能提升工作需要持续不断的努力,另一个即将到来的新突破,是全新的高度优化后的浮点卷积核库XNNPACK[36].在TFLite支持的所有关键浮点卷积模型上的测试结果表明,XNNPACK可以进一步提高执行速度.

Fig. 7 Single thread CPU performance improvement with XNNPACK (Jun 2020)图7 使用XNNPACK后单线程CPU性能提升比例(2020年6月)

以MobileNet V1浮点模型为例,如图7所示,在多种硬件平台上,使用XNNPACK后,单线程CPU性能提升达20%~200%.比如在X86 Windows平台上,有了2倍的提升.我们将在TensorFlow 2.3版本中集成XNNPACK,并计划在2.4版本中对于浮点模型,在所有平台中默认使用XNNPACK.

2.6 硬件加速器代理

我们将TFLite硬件加速接口称delegate[37](代理),它可以把模型计算图的部分或全部代理给另一个执行器,让它们在硬件加速器中执行.在图8示例中,整个计算图一部分在CPU中执行,另一部分子图被代理给硬件加速器执行.

Fig. 8 TFLite hardware delegate example图8 TFLite硬件加速器代理示意

我们不是把算子逐一放到加速器上执行,而是将计算图的整个子图在加速器上运行,这对于GPU或NPU(神经网络加速器)等需要在设备上尽可能多地进行计算,且中间没有CPU互操作的设备而言,是一个巨大的优势.

TFLite有着非常丰富的硬件加速器支持:在Android系统中,支持NNAPI,GPU,EdgeTPU和Hexagon DSP delegates;在iOS系统中,有Metal和CoreML delegates.

1) GPU delegate[38].适用于跨平台.GPU delegate可在Android和iOS上使用,支持32 b和16 b浮点的模型,对许多浮点卷积模型实现了大幅度速度提升,尤其是较大的模型.不过,运行库大小会有小幅度的增加.GPU的驱动可以有不同的后端,Android系统有OpenGL,OpenCL和VulKan后端,以及iOS上基于Metal的后端.我们最近增加了对OpenCL的支持,在多个视觉模型上进行的测试表明,基于OpenCL的GPU delegate的性能提升为CPU上的4~6倍和OpenGL的2倍.图9是三者在Pixel 4上测试后端性能的对照结果.

Fig. 9 TFLite GPU performance图9 TFLite GPU性能

2) Android NNAPI delegate[39].适用于较新的Android设备.NNAPI delegate可用于在具有GPU,DSP和NPU的Android设备上加速模型.它在Android 8.1(API 27+)或更高版本中可用.NNAPI[40]是Android系统中用于加速机器学习的抽象层.在NNAPI和TFLite中,我们会发现高级算子的定义有很多相似之处,因为二者是紧密联系在一起进行开发的.NNAPI是Android平台级的硬件加速器抽象层,我们可以通过TFLite调用它,而供应商比如高通则为NNAPI提供DSP或GPU的驱动程序.在Android Q上,NNAPI真正进入了一个良好的稳定状态,功能和算子方面都正在接近TFLite的同等水平.有越来越多的用户和硬件供应商予以采用,这促进了这些驱动程序的发展.

3) Hexagon delegate[41](DSP).适用于较旧的Android设备.Hexagon DSP是一种微处理器,常见于大量使用高通骁龙SoC的安卓手机.与CPU相比,新的TFLite Hexagon delegate利用DSP实现了MobileNet和InceptionV3等模型,性能大幅提升,提升幅度可达3~25倍,同时CPU和GPU的能效也得到了提升[42].它可以在不完全支持NNAPI的旧版Android设备上使用.

4) Core ML delegate[43].适用于较新的iPhone和iPad.CoreML是在苹果设备上使用的机器学习框架,它还提供了在Neural Engine上运行机器学习模型的API.CoreML delegate允许在CoreML和Neural Engine(如果设备有相关芯片)上运行TFLite模型,以更低的功耗实现更快的推理速度.在含有Neural Engine的iPhone XS以及后续发布的设备上,测试表明各种计算机视觉模型的性能提高了1.3~11倍[44].

图10给出了使用TFLite GPU delegate的Java示例,只需要多加几行代码,非常简单,而其他delegate也类似.值得注意的是,图中的delegate.bindGlBufferToTensor可以使用OpenGL纹理作为计算图输入,这样数据就不用复制,减少在CPU和GPU之间来回传送数据.

Fig. 10 Use TFLite GPU delegate in Java图10 使用TFLite GPU delegate的Java示例

可扩展性是TFLite很重要的设计目标,如果一些硬件后端默认还不支持,完全可以自己定制delegate,只需要:

1) 定义一个新的delegate算子,它负责计算被代理的子图.

2) 创建一个TfLiteDelegate实例,它负责注册新定义的delegate算子,并可以把计算图中的一些子图替换为新的delegate算子.

这样,应用开发者可以利用新的加速器,同时硬件厂商也可以扩展对TFLite的支持,从而触及更多的用户.

我们正在持续改进如何更方便地实现一个新TFLite delegate,以及提供可扩展的工具链来验证测试这一新delegate.目前推荐扩展SimpleDelegate相关APIs[45]来实现,具体可以参考我们提供的简单案例[46]来进一步了解,包括如何集成到我们的测试和验证工具链当中.更多请参考delegate使用开发指南[37].

2.7 TensorFlow Lite微控制器版

1.4节提到,MCU没有操作系统,只有极小内存,低功耗、便宜、无处不在.真正在MCU上部署机器学习模型变得很普及时,智能开始无处不在.TFLite可以支持微控制器MCU,为IoT领域的智能化带来很多想象空间.

TFLite微控制器版(TFLite for Microcontrollers)[47]接受和TFLite同样的FlatBuffers模型格式,让同样的模型可运行在手机和MCU上,减轻了开发者负担.它的核心运行库在Arm Cortex M3上小于16 KB,加上一些关键词检测的算子,总共只需要22 KB.

它支持ARM Cortex-M系列芯片以及其他架构比如ESP32.也可以作为Arduino库,对MBed开发环境也有很好的支持.最近,Cadence宣布旗下的Tensilica HiFi DSP系列也支持TFLite微控制器版[48].

MCU上的模型都非常小,低功耗,只能完成简单的功能,不过可以常驻内存,一直处于等待触发状态.一般常见的用法是级联触发:比如先用MCU的一个模型检测是否有声音,如果是,另一个MCU模型判断这是不是人的声音(排除外界噪音),更进一步才开启设备CPU上更强大的语音识别功能.如果设备需要开启CPU持续监听外界的声音,功耗会很大.

我们开源了很多有意思的样例,比如:

1) 识别若干关键词的语音识别模型,只有20 KB左右.

2) 魔法棒.只需要20 KB左右的模型就可以做姿态检测,这很适合一些使用加速计数据的应用.

3) 人物识别.250 KB的视觉模型识别摄像头是否有人出现.

3 使用TensorFlow Lite的最佳实践

3.1 解决模型转换挑战

目前TFLite约有130个算子,大多同时支持浮点和量化类型,可以在这里找到所有支持的算子[32].TensorFlow有一些语义在TFLite中尚未得到很好的原生支持,同时在模型转换过程中算子进行了优化和融合.总结一下TensorFlow算子和TFLite算子的兼容性:

1) 大多TFLite算子支持float32和量化(int8,uint8),不过很多算子尚不支持bfloat16或string.

2) TFLite目前只支持TensorFlow的某些固定格式,而广播只支持有限的一些ops(比如tf.add,tf.mul,tf.sub,andtf.div,tf.min,tf.max).

3) 一些TensorFlow算子有严格对应的TFLite算子,比如tf.nn.avg_pool,tf.nn.conv2d,tf.nn.depthwise_conv2d,tf.nn.l2_normalize,tf.nn.max_pool,tf.nn.softmax,tf.nn.top_k等,以及tf.matmul,tf.one_hot,tf.reduce_mean,tf.reshape,tf.sigmoid,tf.squeeze,tf.strided_slice,tf.transpose等.详细清单请参考文献[32].

4) 另外很多TensorFlow算子,没有一一对应的TFLite算子,不过TFLite仍可以支持,在模型转换过程得到优化.这些TensorFlow算子可能被删除(比如tf.identity),替换(比如tf.placeholder被换成张量),或者融合为更复杂的操作(比如tf.nn.bias_add).详细清单请参考文献[32].

如遇到模型转换问题,请尝试如下解决方法:

1) 自定义算子.可以自定义新的算子,或者提供定制的更优化的算子.这种方式对于开发者而言最灵活,可控性高.

2) 使用其他等价算子.可以替换为其他TFLite能支持的算子,假设模型更改后精度和速度等方面没有明显损失.

3) Select TF Ops.可以重用TensorFlow算子,缺点是二进制文件大小会有所增加,速度也可能会慢些.

5) 控制流(control flow)支持.最新的基于MLIR的模型转换器已经对此有不错的支持.

6) 动态张量形状(dynamic tensor shape)的支持.目前对动态张量形状支持有限,需要在转化时将输入设置成固定大小,随后在运行时调用ResizeInput.设置输入为固定大小的技巧包括调整图片大小,或者为文字加填充.

一方面,算子在不断增加中,我们也在为TFLite增加动态张量形状等更多新功能的支持.持续关注社区和github更新,可尝试每日更新的代码是否已支持;还可以向社区发issue或者直接贡献代码.

3.2 进行模型优化

压缩模型可以加速模型执行,更好地利用硬件加速器,特别是当一些加速器只能接受全整型的模型时.模型压缩领域是很热的研究话题,而我们则更关注足够简单的开发体验,希望用户只需要几行代码就可以实现.TensorFlow提供了一个模型优化工具包(TensorFlow model optimization toolkit, MOT)[52],简化模型优化过程.MOT被很好地集成到了TFLite工具链中,可在TFLite模型转换时轻松启用.

MOT目前主要支持2类优化:

1) 量化(quantization).降低模型参数精度,比如将浮点转换成整型.

2) 剪枝(pruning).减少模型参数,比如去掉那些不重要的权重(比如为0的权重).

从开发者使用体验来看,又可以分为:

1) 训练后(post-training).比如训练后量化(post-training quantization),不需要改变模型,非常简单.

2) 训练中(training-aware).需要改动模型,在训练的时候优化.

推荐训练后量化(post-training quantization)[53],足够简单,只要能满足需求.

Table 1 TFLite Post-Training Quantization

Fig. 11 Quantization: map 32 b float to 8 b integer图11 量化:将32 b浮点数映射到8 b整数

最初引入的是动态范围量化(dynamic range quantization),它可以动态地对输入和激活函数做量化和去量化,是一种混合量化:当有量化算子支持时,调用量化算子核,否则使用浮点算子核,整个执行过程混合了2类算子.精度上会有所损失,损失多大取决于具体模型.图11展示了如何将32 b浮点数映射到8 b整数.它贵在简单,只需在转化时多加一行代码:converter.optimizations=[tf.lite.Optimize.DEFAULT].

而全整型量化[54]可以同时量化权重和激活函数,这样所有的权重和计算都是整形的,可以进一步改善延迟,减少峰值内存使用量,以及利用仅支持整数的硬件加速器.

为进一步提高精度,我们支持了按轴量化或按通道量化:不会对整个张量使用同一组量化参数,而是对张量中每个通道使用不同的量化参数.为此,需要使用代表性数据集来评估激活函数和输入的动态范围,帮助确定量化参数,这样可以实现与训练时量化大体相当的精确度.对于一个基于图片的模型而言,或许只需要输入30张图片,就能够探索出量化和输出值的空间.

开发者只需创建一个输入数据生成器,并将代表性数据提供给TFLite转换器即可:

converter.representative_dataset=representative_dataset_gen

如有GPU,则考虑使用16 b浮点数,进一步减少精度损失:

converter.target_spec.supported_types=[tf.lite.constants.FLOAT16]

训练时优化,需要对于模型训练过程做一些改变,相对更复杂.为尽量简化用户体验,MOT针对Keras提供了非常便捷的API.目前可以支持:

1) 训练时量化(training-aware quantization).可以量化整个Keras模型和部分Keras模型节点.

2) 剪枝(pruning).去掉一些权重,让模型更稀疏.

比如量化整个Keras模型只需要增加一行:

quantized_model=mot.quantization.keras.quantize_model(keras_model)

MOT在压缩Google的一些关键模型发挥着重要作用,比如1.4节提到的离线语音模型.我们一直致力于开发一些更易用的工具,既简化操作,又能够在模型压缩比和加速性能方面同样优异.比如,之前的训练时量化工具需要在模型中插入伪量化(fake quantization)节点,相对复杂,因此我们发布了新的更易使用的基于Keras训练时量化[55].

我们也在探索更多压缩方法的支持,比如张量压缩算法或者模型蒸馏.

3.3 减少运行库大小

除了压缩模型本身,我们也希望减少运行库大小.一个技巧就是只链接所需要的算子,因为通常一个模型只用到部分算子,我们称之为选择性注册(selective registration).

TFLite有很好的可扩展性,可以自己定义算子和算子解析器.我们在tensorflowlitetools下提供了工具,给定一个模型,可以扫描用到的算子,自动生成一个注册实际算子的代码文件,这样利用自定义的算子解析器,就可以删除不用的算子内核.

3.4 选择正确的模型

选择合适的模型非常关键,这很大程度上取决于应用的需求,比如:

1) 模型大小的限制;

2) 模型延迟的需求;

3) 精度要求;

4) 硬件运行环境.

可以尝试不同的模型,对各个方面做取舍.比如,如果精度要求高而硬件许可,可以选择较复杂的模型;而速度最关键时,选择简单而精度低的模型.

研究领域发布的模型日新月异,开发人员很难时刻跟进,另外,论文上的精度更多是理想情况,是否在真实场景下有稳定表现还需要根据实际场景评估.通常情况下,我们可以从一些经典的针对移动特别优化的在工业界久经考验的模型开始,比如MobileNet.

可以尝试在目标手机上运行一些TFLite的参考应用[56],感受模型实际取得的效果,比如对象追踪应用的流畅程度.

TFHub上提供了一系列的经典模型[57],既有预训练的TF模块,可做迁移学习后转换到TFLite模型,也可以直接下载TFLite模型.机器学习是一个飞速发展的领域,也许每隔几个月就有新的模型刷新记录.我们在努力确保最前沿(SOTA)的模型可以在TFLite上运行.

最近,我们增加了对EfficientNet-Lite(图像分类模型系列)[58],MobileBERT[59]和ALBERT-Lite[60](支持多种NLP任务的轻量级版本BERT)的支持.

EfficientNet-Lite是一种新颖的图像分类模型,可通过减少计算和参数的数量级来实现SOTA的准确性.它针对TFLite量化方式进行了优化,在损失较低精度(几乎可忽略)的同时大大提升了推理速度,并可以运行在CPU,GPU和EdgeTPU上[61].

图12中,在相似的高精度下,EifficientNet-Lite4比Inception V4有显著的速度提升.

Fig. 12 Benchmarked on Pixel 4 CPU with 4 threads (March 2020)图12 在四线程Pixel 4 CPU上做基准测试(2020年3月)

Fig. 13 Benchmarked on Pixel4,Float32 QA model, 4 threads CPU图13 在Pixel 4,Float32问答模型,4线程CPU上评测

MobileBERT[62]和ALBERT-Lite[63]是流行的BERT模型的优化版本,该模型在一系列NLP任务(包括问答、自然语言推断等)上均达到了SOTA的准确度.图13中,MobileBERT的体积是BERT的14,速度是BERT的4倍,同时保持了相近的准确度.ALBERT-Lite的体积甚至更小,仅有BERT的16,也就是MobileBERT体积的23,在速度上ALBERT-Lite稍逊于MobileBERT.我们还在探索量化版本的MobileBERT,它是BERT的116,速度是BERT的8倍,同时也保持了相近的准确度,MLPerf社区正尝试基于此建立移动硬件加速的NLP测试基准.

3.5 性能的最佳实践

TFLite性能测试工具[64]可以找出在特定设备上运行时的性能瓶颈,输出耗时长的算子,还可以插入不同的计算后端,并探索它究竟是如何影响推理性能的.它支持对内部事件(如算子调用)的检测日志记录,可以通过Android的系统跟踪(system tracing)[65]进行追踪.

以下是性能分析工具给出的分析示例,及性能提升的备选解决方案:

1) 如果可用的CPU内核数量少于推理线程的数量,则CPU调度开销可能会导致性能下降.可以在应用程序中重新调度其他大量占用CPU的任务,以避免与模型推断重叠,或尝试调整解释器线程的数量.

2) 如果算子没有完全被代理到GPU上执行,那么模型计算图中的某些部分将在CPU上执行,而不是按预想的在硬件加速器上执行.可以将不支持的算子替换为相似的已支持算子.

尽可能优化模型,比如量化.尽可能地使用硬件加速器,也注意一些性能负担,比如内存零拷贝.

当然,移动平台可能具有多种硬件加速器,而模型的选择也很多(比如模型类型、量化或浮点),组合起来比较复杂.同时这些加速器也可能需要被用作其他用途,比如GPU可能需要预留给视频处理,这带来了优化的复杂度.通常情况下,需要在不同平台上尝试运行应用,统计数据,根据数据进行优化.我们也在探索提供更好的开发者工具支持.

3.6 TensorFlow Lite的使用限制

由于TensorFlow算子丰富,而新模型也不断出现,虽然TFLite覆盖了很多常见的模型,对于普通开发者,模型转换仍是较大的痛点.请参考3.1节部分解决,比如自定义算子,利用Select TF Ops.

TFLite目前对动态张量形状支持有限,需要在转化时将输入设置成固定大小,随后在运行时调用ResizeInput.很多TFLite ops可能还不支持bfloat和string.另外,TFLite目前只支持TensorFlow的某些固定格式,而广播只支持有限的一些ops(比如tf.add,tf.mul,tf.sub,andtf.div,tf.min,tf.max)[32].

硬件加速部分,delegate的算子不支持仍是主要挑战.总体而言,CPU算子有较好的覆盖率,而GPU和DSP等覆盖率相对小一些;而NNAPI的算子则更新速度较慢,且其硬件驱动的性能不一,建议进行实际模型性能测试.如碰到问题,欢迎到社区里反馈.

社区对于个性化的应用和联邦学习(federated learning[66])的兴趣逐步增强,因此对于端侧训练的需求越来越多.目前TFLite对于端侧训练还不支持.我们提供了一个TFLite案例,允许有限的端侧模型迁移,从而让端侧模型个性化,不过使用相对复杂[67].

4 适合初学者的工具

4.1 预训练模型和完整参考示例

我们提供了预训练模型的代码库和实现这些模型的示例应用[56,68],开发者无需编写任何代码即可在实际设备上试用TFLite模型,并可以在github找到完整应用代码,包括模型的前处理和后处理.包括多种案例,比如对象追踪、风格迁移和问题回答,也包括不同平台的例子,比如Android、WiOS、树莓派及MCU.

比如,在手机上实现问题问答[17]是一个很有挑战的问题,我们发布了基于MobileBERT的参考应用,学术界的多个BERT模型可以在这个应用上运行.风格迁移模型[69]启发了Google Arts & Culture的产品新特性[14].

对于视频处理应用,MediaPipe[70]提供了很好的框架,可以和TFLite配合使用.MediaPipe也有丰富应用例子,比如手势追踪.

另外,社区github项目“Awesome TFLite”[71],收集了很多有意思的示例.

4.2 TensorFlow Lite Model Maker

TFLite Model Maker[72]使用迁移学习,可让开发者在自己的数据集上应用最前沿的机器学习模型.它将复杂的机器学习概念封装在直观的API中,无需机器学习专业知识,只需几行代码即可训练最新的图像分类模型:

Fig. 14 TFLite Model Maker example图14 TFLite Model Maker示例

ModelMaker支持TensorFlowHub上的许多最新模型,如3.5节EfficientNet-Lite.如果想获得更高的准确率,仅需修改一行代码,即可切换到不同的模型架构.目前支持图片分类和文本分类任务,我们还在不断扩展其功能.

4.3 TensorFlow Lite Support库

执行模型推理前,通常需要对输入数据做转换,称为前处理,比如图片裁剪和大小挑战.而推理之后,如何解读数据,又需要做后处理.

为简化开发流程,我们提供了TFLite Support[73]库,它提供了一系列工具库.目前支持图片处理库,更多数据类型(比如NLP)和更多平台(比如iOS)的支持正在进展当中.

4.4 模型元数据、自动代码生成和Android Studio工具

直接使用TFLite解释器,它的输入输出都是张量.这带来了2个问题:

1)TFLite模型的使用者将需要确切地知道一个1×224×224×3张量的含义.比如它是位图吗?特别是模型创建团队和使用团队不是同一个时,更容易出问题.

2)对高维数据进行转换时容易出错,例如把Bitmap转换为RGB浮点数组或者ByteArray.

我们希望把模型调用并做前后处理的整个复杂过程都封装到一个简单的API中,开发者只需要调用简单4~5行代码,而且可以直接使用他们熟悉的原生数据,比如Android开发时直接用Bitmap,而不是转化为ByteArray.

为此,TFLite增加了对模型元数据的支持[74],这让模型创建者可以使用类型化的对象来描述模型的输入和输出,方便模型共享和自动化处理.更进一步,提供了一个Android代码生成器[75],给定模型,它可以读取元数据,自动生成封装好的Java类,它可以自动调整图片大小,对其进行归一化且从ByteArray进行转换.

更进一步,我们正将此功能集成到Android Studio的ML model binding工具中[76],只需要把模型导入到Android Studio中就可以生成代码,目前已提供试用版,请参考文献[77].

5 未来的方向

TFLite作为开源项目,为了更好地与社区互动,我们提前公开了开发计划[78],以确保所从事的工作和优先级清晰明了.

性能是我们持续关注的重点.我们正在集成XNNPACK以便于进一步优化浮点模型,还在提升训练后量化的性能,以便加快CPU推理速度.增加更多优化的算子,并让现有的硬件加速器delegate支持更多算子,以便加速更多模型,支持更多的硬件加速器.在有多个异构加速器的情况下,探索如何提供更简易的工具,帮助发挥出其最佳性能.

简化开发者体验很关键.我们会持续不断更新最前沿的端侧模型以及相关的示例,展示更多的可能性.增强TFLite Model Maker,支持更多任务,例如目标检测,支持BERT系列的NLP任务.更丰富的TFLite Support库.扩展模型元数据和代码生成工具,以支持更多用例,以及与Android Studio的无缝集成,进一步简化体验.提供更便捷工具,帮助用户根据模型来减少算子库,压缩运行时间.

进一步让TensorFlow模型到TFLite模型转化更流畅,比如更好地支持RNNLSTM和动态形状、模型转换过程更好的语义匹配(比如只转换部分计算图).长期看来,如何更好地融合TensorFlow和TFLite,共享算子和核心库,同时,也兼具移动平台的优化.

提供更丰富的模型优化工具,比如剪枝、张量压缩和蒸馏技术.

在TFLite微控制器版中,和社区合作支持更多芯片平台、更多算子,同时提供更多适合MCU的样例模型和应用.

更好地支持联邦学习,保护用户隐私.更好地支持端侧训练,支持个性化应用.

致谢感谢TensorFlow团队,特别是Tensor-Flow Lite团队的杰出工作,设计并开源TensorFlow Lite,加速端侧机器学习的发展.感谢Google相关团队持续对TensorFlow Lite的贡献和反馈.感谢广大硬件平台的支持.感谢最广大开发者对TensorFlow社区的贡献,持续推动其发展.