TensorFlow下构建高性能神经网络模型的最佳实践
作者 | 李嘉璇
责编 | 何永灿
随着神经网络算法在图像、语音等领域都大幅度超越传统算法,但在应用到实际项目中却面临两个问题:计算量巨大及模型体积过大,不利于移动端和嵌入式的场景;模型内存占用过大,导致功耗和电量消耗过高。因此,如何对神经网络模型进行优化,在尽可能不损失精度的情况下,减小模型的体积,并且计算量也降低,就是我们将深度学习在更广泛的场景下应用时要解决的问题。
加速神经网络模型计算的方向
在移动端或者嵌入式设备上应用深度学习,有两种方式:一是将模型运行在云端服务器上,向服务器发送请求,接收服务器响应;二是在本地运行模型。
一般来说,采用后者的方式,也就是在PC上训练好一个模型,然后将其放在移动端上进行预测。使用本地运行模型原因在于,首先,向服务端请求数据的方式可行性差。移动端的资源(如网络、CPU、内存资源)是很稀缺的。例如,在网络连接不良或者丢失的情况下,向服务端发送连续数据的代价就变得非常高昂。其次,运行在本地的实时性更好。但问题是,一个模型大小动辄几百兆,且不说把它安装到移动端需要多少网络资源,就是每次预测时需要的内存资源也是很多的。
那么,要在性能相对较弱的移动/嵌入式设备(如没有加速器的ARM CPU)上高效运行一个CNN,应该怎么做呢?这就衍生出了很多加速计算的方向,其中重要的两个方向是对内存空间和速度的优化。采用的方式一是精简模型,既可以节省内存空间,也可以加快计算速度;二是加快框架的执行速度,影响框架执行速度主要有两方面的因素,即模型的复杂度和每一步的计算速度。
精简模型主要是使用更低的权重精度,如量化(quantization)或权重剪枝(weight pruning)。剪枝是指剪小权重的连接,把所有权值连接低于一个阈值的连接从网络里移除。
而加速框架的执行速度一般不会影响模型的参数,是试图优化矩阵之间的通用乘法(GEMM)运算,因此会同时影响卷积层(卷积层的计算是先对数据进行im2col运算,再进行GEMM运算)和全连接层。
模型压缩
模型压缩是指在不丢失有用信息的前提下,缩减参数量以减少存储空间,提高其计算和存储效率,或按照一定的算法对数据进行重新组织,减少数据的冗余和存储的空间的一种技术方法。
目前的压缩方法主要有如下4类:
设计浅层网络。通过设计一个更浅的网络结构来实现和复杂模型相当的效果。但是因为浅层网络的表达能力往往很难和深层网络匹敌,因此一般用在解决简单问题上。
压缩训练好的复杂模型。采用的主要方法有参数稀疏化(剪枝)、参数量化表示(量化),从而达到参数量减少、计算量减少、存储减少的目的。这是目前采用的主流方法,也是本文主要讲述的方法。
多值网络。最为典型就是二值网络、XNOR网络。其主要原理就是采用0和1两个值对网络的输入和权重进行编码,原始网络的卷积操作可以被位运算代替。在减少模型大小的同时,极大提升了模型的计算速度。但是由于二值网络会很大程度降低模型的表达能力。因此也在研究n-bit编码方式。
知识蒸馏(Knowledge Distilling)。采用迁移学习,将复杂模型的输出做为soft target来训练一个简单网络。
下面我们来着重介绍目前应用较多的压缩训练好的复杂模型的方法。
剪枝(Prunes the network)
剪枝就是将网络转化为稀疏网络,即大部分权值都为0,只保留一些重要的连接。如图1所示。
图1 剪枝的过程及剪枝前后的对比:剪枝权重及剪枝下一层神经元
事实上,我们一般是逐层对神经网络进行敏感度分析(sensitive analysis),看哪一部分权重置为0后,对精度的影响较小。然后将权重排序,设置一个置零阈值,将阈值以下的权重置零,保持这些权重不变,继续训练至模型精度恢复;反复进行上述过程,通过增大置零的阈值提高模型中被置零的比例。具体过程如图2所示。
图2 交互式剪枝的过程
剪枝的特点:
通用于各种网络结构与各种任务,且实现简单,性能稳定;
稀疏网络具有更低的功耗,在CPU上使用特定工具时具有更快的计算速度;
剪枝后的稀疏矩阵通常采取特殊的存储方式,例如,常用MKL中的CSR格式。
剪枝的结果:
通过在现有的经典神经网络上做实验,发现压缩倍数在9-12倍之间。如图3所示;
压缩的多是全连接层,CNN层参数少,因此能压缩的倍数也较少;
根据经验,压缩到60%以上模型存储大小,模型大小才会下降比较多。
图3 经典神经网络剪枝前后的参数对比及压缩率
量化(Quantize the weights)
量化(Quantization)又称定点,用更少的数据位宽进行神经网络存储和计算。它的优势在于节省存储,并进行更快地访存和计算。
量化是一个总括术语,用比32位浮点数更少的空间来存储和运行模型,并且TensorFlow量化的实现屏蔽了存储和运行细节。
神经网络训练时要求速度和准确率,训练通常在GPU上进行,所以使用浮点数影响不大。但是在预测阶段,使用浮点数会影响速度。量化可以在加快速度的同时,保持较高的精度。
量化网络的动机主要有两个。最初的动机是减小模型文件的大小。模型文件往往占据很大的磁盘空间,例如,上一节介绍的网络模型,很多模型都接近200MB,模型中存储的是分布在大量层中的权值。在存储模型的时候用8位整数,模型大小可以缩小为原来32位的25%左右。在加载模型后运算时转换回32位浮点数,这样已有的浮点计算代码无需改动即可正常运行。
量化的另一个动机是降低预测过程需要的计算资源。这在嵌入式和移动端非常有意义,能够更快地运行模型,功耗更低。从体系架构的角度来说,8位的访问次数要比32位多,在读取8位整数时只需要32位浮点数的1/4的内存带宽,例如,在32位内存带宽的情况下,8位整数可以一次访问4个,32位浮点数只能1次访问1个。而且使用SIMD指令,可以在一个时钟周期里实现更多的计算。另一方面,8位对嵌入式设备的利用更充分,因为很多嵌入式芯片都是8位、16位的,如单片机、数字信号处理器(DSP芯片),8位可以充分利用这些。
此外,神经网络对于噪声的健壮性很强,因为量化会带来精度损失(这种损失可以认为是一种噪声),并不会危害到整体结果的准确度。
那能否用低精度格式来直接训练呢?答案是,大多数情况下是不能的。因为在训练时,尽管前向传播能够顺利进行,但往往反向传播中需要计算梯度。例如,梯度是0.2,使用浮点数可以很好地表示,而整数就不能很好地表示,这会导致梯度消失。因此需要使用高于8位的值来计算梯度。因此,正如在本文一开始介绍的那样,在移动端训练模型的思路往往是,在PC上正常训练好浮点数模型,然后直接将模型转换成8位,移动端是使用8位的模型来执行预测的过程。
图4展示了不同精度(FP32、FP16、INT8)表示的数据范围。
图4 不同精度(FP32、FP16、INT8)表示的数据范围
量化有2种类型,均为量化和非均匀量化。我们以将32bit浮点表示成3bit定点值为例。如图5,用INT3来近似表示浮点值,取值范围是8种离散取值(-4,-3,…,3)。
如果不论权值的疏密,直接对应,我们称之为“均匀量化”;如果权值密的量化后的范围也较密,权值稀疏的量化后的范围也较稀疏,称之为“非均匀量化”。
图5 均匀量化和非均匀量化的对比图
TensorFlow下的模型压缩工具
我们以TensorFlow下8位精度的存储和计算来说明。
量化示例
举个将GoogleNet模型转换成8位模型的例子,看看模型的大小减小多少,以及用它预测的结果怎么样。
从官方网站上下载训练好的GoogleNet模型,解压后,放在/tmp目录下,然后执行:
bazel build tensorflow/tools/quantization:quantize_graph
bazel-bin/tensorflow/tools/quantization/quantize_graph \
–input=/tmp/classify_image_graph_def.pb \
–output_node_names=”softmax” –output=/tmp/quantized_graph.pb \
–mode=eightbit
生成量化后的模型quantized_graph.pb大小只有23MB,是原来模型classify_image_graph_ def.pb(91MB)的1/4。它的预测效果怎么样呢?执行:
bazel build tensorflow/examples/label_image:label_image
bazel-bin/tensorflow/examples/label_image/label_image \
–image=/tmp/cropped_panda.jpg \
–graph=/tmp/quantized_graph.pb \
–labels=/tmp/imagenet_synset_to_human_label_map.txt \
–input_width=299 \
–input_height=299 \
–input_mean=128 \
–input_std=128 \
–input_layer=”Mul:0” \
–output_layer=”softmax:0”
运行结果如图6所示,可以看出8位模型预测的结果也很好。
图6 生成量化后的模型quantized_graph.pb运行结果
量化过程的实现
TensorFlow的量化是通过将预测的操作转换成等价的8位版本的操作来实现。量化操作过程如图7所示。
图7 TensorFlow下模型量化的过程
图7中左侧是原始的Relu操作,输入和输出均是浮点数。右侧是量化后的Relu操作,先根据输入的浮点数计算最大值和最小值,然后进入量化(Quantize)操作将输入数据转换成8位。一般来讲,在进入量化的Relu(QuantizedRelu)处理后,为了保证输出层的输入数据的准确性,还需要进行反量化(Dequantize)的操作,将权重再转回32位精度,来保证预测的准确性。也就是整个模型的前向传播采用8位段数运行,在最后一层之前加上一个反量化层,把8位转回32位作为输出层的输入。
实际上,我们会在每个量化操作(如QuantizedMatMul、QuantizedRelu等)的后面执行反量化操作(Dequantize),如图8左侧所示在QuantizedMatMul后执行反量化和量化操作可以相互抵消。因此,在输出层之前做一次反量化操作就可以了。
图8 量化操作和反量化操作相互抵消
量化数据的表示
将浮点数转换为8位的表示实际上是一个压缩问题。权重和经过激活函数处理过的上一层的输出(也就是下一层的输入)实际上是分布在一个范围内的值。量化的过程一般是找出最大值和最小值后,将分布在其中的浮点数认为是线性分布,做线性扩展。因此,假设最小值是-10.0f,最大值是30.0f,那量化后的结果如表1所示。
表1 量化数据的表示
ResNet50上的模型压缩实验
笔者在ResNet50-v1上,采用官方GitHub上提供的模型作为Baseline,在ImageNet测试集5万张图片上进行测试。结果如图9。
图9 ResNet50网络量化前后的精度对比
在均匀量化的过程中,首先是仅仅对权重进行量化,得到精度为72.8%。随后,分别用模型对测试集的10张、1000张图片的范围进行提前计算最值(Max和Min),并进行存储,得到的精度分别为72.9%和73.1%。
从量化前后的可视化模型对比,也可以看成量化对模型做了哪些操作。图10是未经量化的原始模型。
图10 ResNet50原始网络的节点结构
图11仅仅对权重进行量化,没有计算输入图片的最值范围的可视化模型。可以看出原本的Conv2D等节点都转换为QuantizedConv2D的对应节点。并且在进行QuantizedConv2D操作后,得到INT32类型的记过,需要对操作的结果转换为8位(ReQuantize操作),而转换的过程需要知道INT32结果的最值范围,因此也加入了ReQuantizationRange节点。
图11 仅量化权重,在 Conv2D节点计算后,需要ReQuantizationRange来得到计算后的范围
如果已经预先使用10张或者1000张图片计算了每一个Conv2D等操作之后需要计算的范围,则ReQuantizationRange的计算过程就可以省去,直接从存储的计算好最值文件中读取。如图12所示。
图12 事先计算好1000张图片的范围,可以省去ReQuantizationRange节点
其他建议
在性能受限环境下,对开发者还有没有技术和工程实现方面的其他建议呢?
设计小模型
可以将模型大小做为约束,在模型结构设计和选择时便加以考虑。例如,对于全连接,使用 bottleneck是一个有效的手段。
例如,我们使用TensorFlow官方网站提供的预训练好的Inception V3模型在此花卉数据集上进行训练。在项目根目录下执行:
python tensorflow/examples/image_retraining/retrain.py \
–bottleneck_dir=/tmp/bottlenecks/ \
–how_many_training_steps 10 \
–model_dir=/tmp/inception \
–output_graph=/tmp/retrained_graph.pb \
–output_labels=/tmp/retrained_labels.txt \
–image_dir /tmp/flower_photos
训练完成后,可以在/tmp下看到生成的模型文件retrained_graph.pb(大小为83M)和标签文件retrained_labels.txt。
我们看到,上述命令行中存储和使用了“瓶颈”(bottlenecks)文件。瓶颈是用于描述实际进行分类的最终输出层之前的层(倒数第二层)的非正式术语。倒数第二层已经被训练得很好,因此瓶颈值会是一个有意义且紧凑的图像摘要,并且包含足够的信息使分类器做出选择。因此,在第一次训练的过程中,retrain.py文件的代码会先分析所有的图片,计算每张图片的瓶颈值并存储下来。因为每张图片在训练的过程中会被使用多次,因此在下一次使用的过程中,可以不必重复计算。这里用tulips/9976515506_d496c5e72c.jpg为例,生成的瓶颈文件为tulips/9976515506_d496c5e72c.jpg.txt,内容如图13所示。
图13 bottleneck文件的内容
再如,Highway、ResNet、DenseNet这些带有skip connection结构的模型,也可以用来作为设计窄而深网络的参考,从而减少模型整体参数量和计算量。
还如,SqueezeNet网络结构中通过引入1x1的小卷积核、减少feature map数量等方法,最终将模型大小压缩在1M以内,分类精度与AlexNet相当,而模型大小仅是AlexNet的1/50。
模型小型化
一般采用知识蒸馏。在利用深度神经网络解决问题时,人们常常倾向于设计更复杂的网络,来得到更优的性能。蒸馏模型是采用是迁移学习,通过采用预先训练好的复杂模型(Teacher model)的输出作为监督信号去训练另外一个简单的网络,得到的简单的网络称之为Student model。实验表明,蒸馏模型的方法在MNIST及声学建模等任务上都有着很好的表现。
总结
随着深度学习模型在嵌入式端的应用越来越丰富,例如安防、工业物联网、智能机器人等设备,需要解决图像、语音场景下深度学习的加速问题,减小模型大小及计算量,构建高性能神经网络模型。
本文重点讲解模型压缩和剪枝方法带来的模型大小和计算量的下降,并且能使精度维持在较高水平。除此之外,剪枝的敏感度分析和重新训练(Retrain)也有很多不同的手段;量化也可以在更低精度(5bit、6bit、甚至二值网络)上尝试,笔者也正在进行相关实验,期待和读者一起探讨。
作者简介:李嘉璇,《TensorFlow技术解析与实战》作者,多个技术大会深度学习讲师。有处理图像、社交文本数据情感分析、数据挖掘等实战经验。曾任职百度研发工程师,目前研究构建高性能的神经网络模型及TensorFlow下的压缩工具链,包括模型量化、剪枝。
本文为《程序员》原创文章,未经允许不得转载,更多精彩文章请点击「阅读原文」订阅《程序员》
精选推荐
经验 | Pytorch还是Tensorflow?英伟达工程师帮你总结了
图解TensorFlow架构与设计
视频教程| 使用TensorFlow搭建一个识别手写数字的分类器