从TensorFlow Lite源码入门CNN量化
前言
从2012年AlexNet夺得当年ImageNet冠军开始,深度学习就开始呈爆破式发展,这个过程中迭代出很多优秀的深度学习代码框架,随着技术逐渐成熟,工业界对技术落地的需求越来越迫切,近几年的论文也慢慢分为两大派系,一派主要研究如何提高performance,另一派研究如何提高computational efficiency,当下已经诞生了很多模型压缩和加速方法,这些方法降低了服务器的运行负担甚至可以让高performance网络运行在移动设备上,具体可参考论文的Introduction,论文中提到的量化方法使得模型可以在移动设备上实时运行。
学术界对精度、自由度的要求和工业界对速度、精简度的要求形成了反差,这就使得越来越多的框架开始把training和inference分开,各公司都开始针对移动设备纷纷推出高性能inference库,Nvidia有TensorRT,Intel有OpenVINO,Facebook有Caffe2go,腾讯有ncnn,百度有paddle-mobile,Google有TensorFlow Lite,大厂的思路总是殊途同归,设计思路都是在training和inference中间加一层optimization,使得学术界和工业界完美对接,每当学术界出来一篇paper,我们可以用诸如TensorFlow、Caffe此类框架快速实现,然后把training得到的模型经optimization转换后直接用在工业界。
本文以TensorFlow为例,探究量化的实现方法。
TensorRT
OpenVINO
TFLite自17年11月release出来后,到现在为止已经有classify、detect、stylize、speech四个demo了,抽两天试了下demo、看了下源码,由于官网的文档暂时不完善,先记录一下。
量化的原因
先来认识两个概念
FLOPS: floating point operations per second (FLOPS, flops or flop/s) is a measure of computer performance. 详见wiki
FLOPs: floating point operations.
由此可见FLOPS衡量硬件,FLOPs衡量模型
FLOPS
CPU的FLOPS的计算公式比较明确,如下
No. of Cores * Average frequency * No. of SIMD Units * (No. of mul-add units*2 + No. of mul units)
可以参考
GPU的FLOPS比较迷,不同架构计算方式不一样,不过也都大同小异,Jetson TX2 FLOPS计算方法如下
No. of SMs * No. of FP32 units per SM * 2 * clock rate
可以算得TX2的FLOPS为2*128*2*1.3
= 665.6 FP32 GFLOPS,也就是1331.2 FP16 FLOPS,20.8 FP64 FLOPS。
FLOPs
卷基层和全连接层的FLOPS计算方法见NVIDIA paper的APPENDIX。
经典的目标检测算法YOLOv3-416的模型复杂度为65.86 FLOPs(见YOLO),这样可以计算一下,在TX2上跑YOLOv3-416的模型大概可以跑到665.6/65.86 = 10.1 FPS,当然这只是个理论值,因为inference前还要对数据进行处理,其实darknet中前期的图像处理占用了比较长的时间。
一般情况下CPU的浮点运算能力都比较差,想在ARM设备上实时运行CNN必须针对ARM的结构进行优化。
从ARMv7开始,arm提供了NEON结构,该结构提供了32个64位D寄存器(ARMv8有32个128位寄存器),这32个64位寄存器也可以看做16个128位Q寄存器,幸运的是NEON提供了8bit intrinsics,如果我们能把FP32降低为uint8,gemm就可以用NEON进行快速计算了。
有关NEON的具体指令和编程相关问题可以参考NEON Programmer's Guide。
TFLite底层就是利用了NEON进行gemm,我们下面慢慢分析。
量化原理
参考论文Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference和gemmlowp的文档,gemmlowp的文档写的更详细一些,举了一些例子,介绍了一些程序上的优化方法。
TFLite代码构架
TF Lite的构架写的也很elegant,大概看了下,主要分下面几个level
1. User
这块都很简单,唯一比较麻烦的是需要了解一些Android的知识和Java的知识,最关键函数就是runForMultipleInputsOutputs,分类、检测等inference功能的入口函数。
2. Interpreter
该层所有函数基本都在org/tensorflow/lite目录下,有三个比较关键的函数:
- createModel();createModel用来反序列化二进制文件生成Model。
- CreateOpResolver();将所有Operator都注册成TfLiteRegistration,这样所有的Operator都可以直接通过TfLiteRegistration来调用了,可以认为每个Operator都是TfLiteRegistration的子类。对比下caffe,每个Operator对应Caffe中的某个layer,TfLiteRegistration对应Caffe中的父类Layer,比如Operator_RELU对应Caffe中的relu_layer,对比只是为了容易理解,和Caffe只是实现方式不同而已,Caffe用工厂类来实现,TFLite用函数指针来实现。
- createInterpreter;用来解析Model中的SubGraph、Tensor和Operator,并把Operator转换成TfLiteRegistration类型,详见Module层。这个函数有一个比较重要的操作,就是创建了一个interpreter类,这个类中有一个成员变量是TfLiteContext,context是TFLite比较重要的一个类,主要决定了是否进行量化、要用什么Kernel进行计算等。
- run;调用Interpreter类中的Invoke()函数,依次调用各个TfLiteRegistration中的invoke函数(也就是各个Operator中的Eval函数)实现inference。
这几个函数在native文件中找到对应的c++实现,通过完成调用Module层的函数来实现各功能。
3. Module
TFLite网络架构的序列化、反序列化是用FlatBuffers实现的,该层代码的阅读可以从schema.fbs开始看,主要由Model、SubGraph、Tensor、Operator几个部分构成,很好理解,和Caffe对比一下,SubGraph对应Net、Tensor对应Blob,Operator对应layer,Model可以有多个SubGraph,不过当前只支持一个SubGraph。
4. Kernel
Kernel可以调用eigen和gemmlowp,主要实现gemm操作,gemmlowp也是由Context来进行计算的,在TFLite中封装成了RefCountedGemmContext类,继承自TfLiteExternalContext类,详见http://gemm_support.cc,调用TfLiteRegistration的invoke就会调用对应Operator的Eval函数,gemmlowp的入口在optimized_ops.h中,也就是gemmlowp::GemmWithOutputPipeline,gemmlowp的具体应用可以参考quantization_example。
gemmlowp
之前没接触过底层的程序,所以这块儿的代码看起来很吃力,我先汇总一下基础知识
- Why GEMM is at the heart of deep learning,可以帮助理解下gemm的原理以及其在CNN中的重要性。
2. 告别四大误区 谈缓存对CPU的性能影响,帮助小白理解CPU命中率、L1与L2 Cache等概念。
3. Gallery of Processor Cache Effects,文章有中文翻译7个示例科普CPU Cache 用7个程序来帮助理解cache line的工作原理。
有了这些基础之后,代码读起来才会有感觉,不然一脸懵逼。
gemmlowp的文档写的非常好,比较难理解的部分都在doc中给出了详细的解释,我按照我的理解总结一下
代码执行过程分pack、compute product using kernel、unpack三个步骤。
1. pack
pack是三个步骤中最难理解的,涉及到L1、L2缓存以及CPU的寄存器等相关知识,再理解了上述基础知识2和3之后会理解pack就相对容易些,pack也分了三个步骤,packL2->packL1->packKernel,先把原始数据打包成L2 cache的结构,然后根据L1 cache的大小从packedL2中取出一部分数据作为packL1,然后把packL1的数据Prefetch到cache(我查了以下,现在的编译器优化都比较好,Prefetch貌似是没有太多作用,但我没做过实验,不敢下结论),再根据寄存器的大小、数量从packedL1中拿取一部分数据用于gemm,代码中比较关键的就是pack.h中的pos_指针的移动。
2. compute product using kernel
这部分基本都是汇编语言实现的,根据不同平台的指令集使用不同的计算方式,NEON的实现方式可参考kernel_neon.h。
3. unpack
这部分的理论可以参考low-precision.md,写的已经很清晰了,不再赘述。
看一下我测试效果
浮点计算效果
量化的计算结果
单从这个简单的例子看,8bit和FP32误差并没有那么大,但我也没有在CNN实例中测试,不知道效果怎样。
Demo
先看下官方给的测试参数
备注:Pixel用的芯片是骁龙835(在此吐槽下Google的手机,当年买了Nexus 6p,用了不到一年就扔了,在国内是真没法用)
Demo也是很简单就可以运行出来(手机处理器为A57)
欢迎研究CNN量化、压缩、剪枝的同学和我一起探讨交流