基于优化神经网络的遥感影像语义分割算法

王晓锋 张猷 夏严峰 王森 王辉

摘要：利用机器学习技术可以自动化地识别地表覆盖物，减轻人力成本和减少人为误差，加快测绘、国防、农业和防灾减灾等各个领域的数字化转型。遥感影像地表覆盖物的分类算法已经得到了一定的发展，但具体到应用场景中，受制于模型效率，自动化水平仍不高。为了保证训练产生的机器学习模型可以稳定、快速地服务于上述的领域，本文针对机器学习模型在训练、测试阶段的优化进行了探索，从而保证模型可更高效稳定地在有限的硬件资源上运行。通过本文介绍的优化方法，模型的训练速度得到了提升，模型占用的空间得到了减小，模型的推理速度得到了提高，并且保证了模型的精度几乎不受影响。

关键词：神经网络压缩;神经网络优化;深度神经网络;图像语义分割;遥感影像

中图分类号：TP391      文獻标识码：A

文章编号：1009-3044（2022）17-0070-03

1 概述

航拍图像和遥感图像识别长期以来得到了测绘、国防、农业等领域的广泛应用，随着深度学习和卷积神经网络（Convolutional Neural Networks，CNN）的高速发展，地表覆盖物分类等基于遥感影像的自动化分析技术也得到了一定程度的发展。

在过去的10年内，视觉模型的发展主要有两个主流方向，分别是增加模型对于图像的理解能力和降低模型理解图像的算力消耗。通过增加模型复杂度提高精度已经得到了充分研究，但如何在不损失太大精度的前提下，降低模型的大小和算力消耗仍有待探索。为了使深度学习模型得到更广泛的应用，尤其是在移动端设备上应用，这样的探索是必须的。

目前，学术界已有一些模型优化的成果，比如MobileNetV2[1]和ShuffleNetV2[2]。从优化的不同时间阶段，模型优化可以分为训练阶段和推理阶段。在模型训练阶段，可以更改模型权重更新的过程，从而提高训练效率;在模型推理阶段，可以将训练好的模型进行编译优化，得到一个体积更小、运行更快的部署模型。

（1）训练阶段的优化

l空间可分离卷积：在不考虑通道维度的情况下，通过将卷积核分解降低运算量，但并不是任何卷积核都可以分解为两个向量，因此不具备通用性。而且，计算矩阵分解还需要消耗一定的算力。

l深度可分离卷积：在考虑通道维度的情况下，将矩阵连乘的过程优化为乘法的加法。需要注意的是，深度可分离卷积是减少卷积的参数量的近似算法，和朴素的卷积并不是严格等价的。

l量化：在不过分影响模型效果的前提下，将部分32位浮点数精度的数值降低为16位浮点数和8位整数，这样可以减少模型体积、加快训练的速度。

l剪枝：在不过分影响模型效果的前提下，将模型的部分结构删除，从而减少模型体积、加快推理速度。

l聚类：通过共享权重的方式来减少权重的数量，这样可以减少模型体积、加快模型收敛速度。在样本稀疏或者样本分布不均衡的部分场景下，可以提升模型的效果。聚类对于模型效果和收敛速度的提升都是鉴于当前的深度学习模型的参数数量很大、拟合能力很强这一事实。

（2）推理阶段的优化

l常量折叠：在编译阶段将一些常量表达式的值计算出来，然后写入二进制文件中（在模型编译优化的场景下，就是优化后的模型），同时也是许多现代编译器使用的编译优化技巧。这样可以通过编译时的单次计算换取运算时的若干次重复运算。并且，由于把常量表达式替换为常量，也减少了不少空间的存储，因此也可以减少模型的体积。

l减少死区代码：经过模型的优化之后，可能有许多的计算步骤被跳过了，因此会出现大量的不会被触达的指令区域。减少死区代码的方法是，通过遍历所有指令来将这些代码移除掉，达到减少模型体积的效果。

l算子融合：由于现代的机器学习建模框架都是以微型算子与图的形式来表示一个模型。适当地将一些常见的算子拓扑的序列融合成一个算子可以极大地减少算子的内存占用和算子之间的数据传输压力，这样就会达到减少模型大小和加快模型计算速度的目的。

训练阶段广泛被使用的框架有TensorFlow[3]和PyTorch[4]等，推理阶段有ARM NN和TensorRT等。虽然这些框架都具有一定的跨平台特性，但要想做到对所有的目标设备（CPU、GPU、NPU、TPU）都能实现很好的推理性能是很难的。这是因为，各种硬件设备的计算特性有所差异，所以框架之间并不存在强通用性。

本文介绍了一些重要的深度学习模型优化方法，并基于遥感影像语义分割问题进行了实验。实验结果显示，本文介绍的方法在不损失太多精度的前提下，对于缩小模型大小、提高训练和推理效率有着很好的效果，有利于遥感影像语义分割算法在移动端的部署与应用。

2 方法

本文介绍神经网络训练和推理阶段的若干优化方法，具体的流程如图1所示。在训练阶段，引入混合精度、权重剪枝、权重聚类获得模型1;在推理阶段，通过编译优化获得模型2，即最终模型。

2.1 混合精度

如今，大多数深度学习模型使用的数据精度是32位浮点数，占用4字节内存。但是，使用16位浮点数可以达到更快的运算速度和更小的内存占用，因此产生了混合精度技术。混合精度是训练时在模型中同时使用16位和32位浮点类型，从而加快运行速度、减少内存使用的一种训练方法。让模型的某些部分保持使用32位类型的目的是，既能保持数值稳定性，又能缩短模型的单步用时，在提高训练效率的同时仍可以获得同等的模型精度。

在混合精度训练时，由于权重、激活值、梯度等都使用16位存储，会更容易产生舍入误差和数据溢出的问题。为了使模型的准确性不损失过多，可以使用权重备份和损失缩放的方法[5]。权重备份是指，在优化步骤中，维护、更新一份32位的权重。在每次迭代中，32位权重被转为16位，用于前向传播和反向传播。权重备份主要防范的问题是，当产生过小的梯度，乘学习率后很容易下溢为0，对模型准确性产生不利的影响。使用权重备份后，可以通过备份的32位权重恢复准确性。损失缩放的方法启发于网络训练期间对数值范围的观察，大部分16位可表示范围未被使用，而许多低于可表示最小值的值变成了0。如果将这些数值增大到2的幂次倍，可以避免下溢产生的误差。最有效的方法是，将前向传播后得到的损失值缩放，再进行反向传播，根据链式法则，后续的梯度也会被缩放同样的倍数。在选取缩放倍数的时候，只要不导致溢出，缩放因子越大越好。如果产生溢出，将使权重或梯度产生无穷大或NaN，而这是无法恢复的。

2.2 权重剪枝

权重剪枝通过删除神经网络层之间的连接，从而减少参数，进而减少计算中涉及的参数和运算量，有利于模型压缩，对于模型的存储与传输来说都有好处。权重剪枝方法可以消除权重张量中作用不大的值，这一步通常会用到某些启发式方法，比如计算权重张量的平均激活值来对权重张量中的值进行优先级排序，找到一定范围内的低优先级的权重，再通过将神经网络中这些参数的值设置为0，以消除神经网络各层之间不必要的连接。由于在权重剪枝过程中，如果一次性剪枝掉太多张量中的值，可能会对神经网络造成不可逆的损伤，因此，权重剪枝方法常常是迭代式进行的，即训练与剪枝过程的交替重复。稀疏张量可以实现很好的压缩效果，因此很适合进行压缩。权重剪枝过程之后，由于张量中存在大量为0的参数，所以可以很好地进行模型压缩，以达到降低模型大小的目的。

下面叙述一个具体的剪枝方法。对于选择修剪的每一层，添加一个二进制掩码变量，与每层的权重形状相同，用于确定哪些权重参与运算。前向传播时，对权重按绝对值进行排序，绝对值最小的权重对应的掩码被设为0，直到达到所需的稀疏度。反向传播过程中，在前向传播中被屏蔽的权重不会得到更新。需要注意的是，为了避免对模型的准确性产生太大影响，设定的稀疏度需要经过多次调整。学习率大小与其调整方式是另一个对模型精度影响很大的地方。如果学习率会下调很多，在学习率大幅下降、网络已经被剪枝后，网络将难以从剪枝带来的准确率下降中恢复。如果初始学习率过高，可能意味着当权重尚未收敛到一个好的解时修剪权重。这里描述的剪枝方法不依赖于任何特定网络的属性，具有高度的泛化性。

2.3 权重聚类

权重聚类通过用相同的值替换权重中的相似权重来减小模型的大小。通过在模型的训练权重上运行聚类算法，可以找到这些相似的权重，并把这些相似的权重用同样的值進行代替。权重聚类在减少模型存储空间和传输大小方面具有直接优势，因为具有共享参数模型的压缩率要比不具有共享参数的模型高得多。这样的模型压缩方法类似于权重剪枝，不同之处在于，权重聚类是通过增加共享权重的参数个数来实现的，而剪枝是通过将低于某个阈值的权重设置为0来实现的。由于权重聚类会将模型中所有参数替换为给定个数的共享权重值，所以在应用权重聚类方法之前，通常需要对模型进行一些预训练，使得模型的参数更新到一个可接受的水平之后，再使用权重聚类对参数进行微调。将模型进行权重聚类后，可以应用常用的压缩算法进行压缩。由于权重聚类后的模型中存在大量的共享参数，可以实现更高的压缩率。

在具体实现过程中，由于聚类中心对于聚类效果有影响，需要选择合适的权重初始化方法，根据文献[6]，线性初始化是最优的方法。计算相似权重的常用方法是K-means，并用欧氏距离度量权重和聚类中心的距离。聚类完成后的聚类中心就是共享权重。

2.4 编译优化

编译优化是将编译前端生成的模型在编译器中端通过计算图的优化方法对模型进行优化。需要注意的是，针对不同的目标设备，使用的优化方法以及优化策略是有所不同的，以确保生成优化后的计算图能够在目标设备上表现出最佳性能。

编译优化的方法主要有：运算符合并，将多个小运算（加、和）合并为一个运算;常量折叠，可以预计算能够确定的部分计算图;静态内存规划，可以预分配临时变量占用的空间，避免多次内存申请和释放;数据分布变换，可以将数据分布变换为有利于进行高效后端运算的分布，比如将行主序或列主序的矩阵变换为更适合目标硬件的存储形式。

3 实验与分析

3.1 实现细节

本文以遥感影像地表覆盖物分类为例对所介绍的方法进行验证。遥感影像数据的来源为上海市某地区的航拍图，并按照建筑、水系、植逐像素地进行了人工标注。完成语义分割任务的基准深度学习网络是深度为4的U-Net[7]，并融合了MobileNetV2[1]，使用TensorFlow实现。输入网络的图片大小均为1024*1024。训练批次大小设为4，迭代20代，损失函数使用Lovasz Loss[8]。

混合精度使用TensorFlow的Keras API实现，权重剪枝和权重聚类使用TensorFlow提供的工具包Pruning API和Weight Clustering API实现，编译优化使用TVM[9]实现。在模型进行储存时，权重仍然按照与权重剪枝或权重聚类前的模型一样的Tensor格式进行存储，而并非会因为权重矩阵变得更为稀疏而采用稀疏矩阵的方式进行存储。所以，为了减少模型权重矩阵中的冗余数据的存储，在存储模型时采用gzip算法进行压缩。

3.2 实验结果

混合精度实验的结果如表1，有预训练时混合精度会缩短训练时长并降低Loss，无预训练时则相反。使用混合精度后，缩短训练时长的表现并不稳定，并且会在一定程度上影响训练结果，但内存占用被大幅减少了。

权重剪枝与权重聚类的实验结果如表2，其中耗时表示一张图片的推理耗时。可以看出，权重剪枝与权重聚类对于推理耗时的影响非常有限，但是能大大缩减模型大小，而且对输出准确度没有负面影响。之所以对推理耗时几乎没有影响，是因为尽管这两种策略会提高矩阵的稀疏度，但是由于TensorFlow默认存储的方式仍然是Dense Tensor，所以在进行计算时，没有对稀疏张量提供计算加速的支持，仍然采用默认的计算方法。

权重剪枝与权重聚类后对一张图的预测结果如图2所示，从左向右分别为：原图、原预测图、应用权重剪枝的预测图、应用权重聚类的预测图。可以看出，权重剪枝与权重聚类对于预测结果没有太大的影响，但权重剪枝会造成比权重聚类更多的误分类。

使用TVM进行编译优化实验的结果如表3，对Conv2D算子进行图优化。可以看出，经过多轮优化，推理时间逐步减少。

4 结论与展望

本文针对遥感图像语义分割在移动平台应用的问题，研究了神经网络优化算法。本文探索了混合精度、权重剪枝、权重聚类、编译优化的方法，在开源平台上进行了实验，分析了各种方法的优势与不足。混合精度能够大幅降低内存使用，但对于训练时间和训练效果的影响不可预料;权重剪枝和权重聚类能大幅度减小存储模型的大小，对推理时间的影响几乎可以忽略不计;编译优化和平台具有很高的相关性，难以部署。为了进一步探索优化算法的适用性，可以在不同的网络模型上进行实验，在不同平台上测试性能表现。

参考文献：

[1] Sandler M，Howard A，Zhu M L，et al.MobileNetV2：inverted residuals and linear bottlenecks[C]//2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.Salt Lake City，UT.IEEE，2018： 4510-4520.

[2] Ma N N，Zhang X Y，Zheng H T，et al.ShuffleNet V2：practical guidelines for efficient CNN architecture design[C]//Computer Vision – ECCV 2018，2018： 116-131.

[3] Abadi M， Agarwal A， Barham P， et al. Tensorflow： Large-scale machine learning on heterogeneous distributed systems[J]. arXiv preprint arXiv：1603.04467， 2016.

[4] Paszke A， Gross S， Massa F， et al. Pytorch： An imperative style， high-performance deep learning library[J]. Advances in neural information processing systems， 2019（32）： 8026-8037.

[5] Micikevicius P， Narang S， Alben J， et al. Mixed precision training[J]. arXiv preprint arXiv：1710.03740， 2017.

[6] Han S， Mao H， Dally W J. Deep compression： Compressing deep neural networks with pruning， trained quantization and huffman coding[J]. arXiv preprint arXiv：1510.00149， 2015.

[7] Ronneberger O，Fischer P，Brox T.U-net：convolutional networks for biomedical image segmentation[C]//Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015，2015：234-241.

[8] Berman M，Triki A R，Blaschko M B.The lovasz-softmax loss：a tractable surrogate for the optimization of the intersection-over-union measure in neural networks[C]//2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.Salt Lake City，UT，USA.IEEE，2018：4413-4421.

[9] Chen T，Moreau T，Jiang Z，et al.TVM：An automated end-to-end optimizing compiler for deep learning[C]//Proceedings of the 13th USENIX Symposium on Operating Systems Design and Implementation（OSDI 8）.Carlsbad，USA，2018：578-594.

收稿日期：2022-02-10

作者簡介：王晓锋（1982—），男，上海人，高级工程师，硕士，主要研究方向为土木工程;张猷（1982—），男，上海人，高级工程师，本科，主要研究方向为输变电工程建设管理;夏严峰（1990—）男，上海人，工程师，硕士，主要研究方向为输变电项目管理;王森（1999—），男，黑龙江庆安人，本科，主研方向为图像分割;王辉（1997—），男，安徽庐江人，硕士在读，主要研究方向为图像分割。