从TensorFlow Lite源码入门CNN量化

从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

注：公式见Jetson TX2 Performance

可以算得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

之前没接触过底层的程序，所以这块儿的代码看起来很吃力，我先汇总一下基础知识

1. 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量化、压缩、剪枝的同学和我一起探讨交流