基于迭代剪枝VGGNet的火星图像分类

火星是21世纪深空探测的主要目标。为了进一步探索火星，人类相继发射Spirit、Opportunity和好奇号着陆火星表面。这些火星车均携带了全景相机。Spirit在运行的6年时间里，分析了Gusev Crater底部的玄武岩、火山岩和火山碎屑岩［1］。Opportunity在运行期间穿越40余公里，探索了火山口中的沉积岩［2］和粗玄武岩［3］。寻找古代沉积岩，已成为火星车的重点工作。这是因为古代沉积岩是液态水存在的证据，也是寻找外星生命的线索［4］。研究岩石还将为地形的物理性质提供语义描述，有助于火星车避障［5］。火星环绕器需要作为中继通信卫星，为火星车建立稳定的中继通信链路，才能传回火星岩石图像供科学研究分析［6］。这就要求火星车对拍摄的火星图像具有较强的分类能力，传回重要的火星图像，以节约火星车的能量。因此，提高火星图像的分类能力具有重要的实用价值和必要性。

近年来，鉴于深度学习算法在图像分类任务上的优良性能以及有用特征的推断能力，使得它们在火星探测任务中具有广泛的应用前景。Anirudh提出了一种深度学习算法MRSCAtt，该算法联合使用空间和通道注意机制来弥补火星图像在色调、质量、照明和清晰度方面与地球图像之间的差异，从而进行精确的图像分类［7］。Li利用VGG-16和迁移学习进行火星图像分类，在MSL Analyst火星岩石图像库上取得了较好的分类效果［5］。Shang提出基于模糊粗糙特征选择与支持向量机结合的图像分类器，用于解决类别多、数量大等问题［8］。Wagstaff为帮助科学家找到感兴趣图像，开发了基于内容搜索的图像分类器［9］。Lu在美国国家航空航天局的行星数据系统图像制图节点中部署了多种经典分类器［10］。柳思聪研制了“天问一号”着陆区域的高分辨率三维地形图，并通过深度卷积神经网络对着陆区地形地貌进行了分类［11］。

虽然这些深度学习图像分类算法应用于真实火星图像时可以取得较好效果［12］，但是无法部署到火星车上。这是由于受到两个限制：卷积神经网络的复杂结构和内存容量。火星车的硬件资源宝贵，图像分类任务不能占用所有计算资源。

为了解决这个问题，本文利用迭代裁剪方式，对VGGNet［13］进行模型压缩。在保持VGGNet高分类准确率的基础上，减小VGGNet占用的内存资源和计算资源。需要注意的是，需根据剪枝后的权重参数重新训练VGGNet。对于裁剪较敏感的层，使用小的裁剪力度或者不裁剪。采用迭代裁剪方式对模型进行压缩有助于减少图像分类器的内存占用量。该方法具体如下；首先，采用迁移学习训练网络的连通性，以便评估神经元的重要性；其次，通过迭代剪枝方法修剪不重要的神经元，以便将全连接层的参数量和内存占用量减少；最后，采用K-means++聚类实现权重参数的量化，利用霍夫曼编码压缩迭代剪枝与量化后的VGGNet权重参数，达到减少存储量和浮点数运算量的作用。此外，本文通过5种数据增强方法进行数据扩充，目的是解决类别不平衡的问题。为了验证压缩后的模型有效性，将其与轻量级ShuffleNet［14］、MobileNet［15］和EfficientNet［16］模型进行了比较。实验结果表明，所提模型具有更好的性能。

2 数据集及数据处理

2.1　数据集

Mars32K数据集为好奇号火星车MAHLI相机在2012年8月至2018年11月期间采集的图像。该数据集给出了火星的地理和地质特征，如山脉、山谷、陨石坑、沙丘和岩石地形。本文对Mars32K火星表面数据集进行类别划分，表1给出了岩石类别和对应标签，图1给出了7种火星图像示例。玄武岩高度氧化，呈弱碱性，通常二氧化硅不饱和［17-18］。层状沉积岩是一种岩石露头，显示出层状特征［19］，其分层称为“交叉层理”，表明砂岩是由风沉积而成的迁移沙丘［20］。细粒岩是火星地表暴露最多的砾石，其质地呈斑状或隐晶质，有时呈水泡状［21］。砾岩在不同深度的分离结晶含水量不同，产生了粗面岩、粗面玄武岩、夏威夷岩以及堆积的辉石岩、沙丘岩。粗玄武岩的来源尚未明确，可能是从陨石坑壁脱落或火山喷射物［22］。机车和沙顾名思义，不再赘述。

表1  图像类别及其标签

Tab.1  Image categories and their labels

类别描述 标签

玄武岩 G1

砾岩 G2

机车 G3

细粒岩 G4

层状沉积岩 G5

沙 G6

粗玄武岩 G7

下载: 导出CSV

图1  7种图像类别示例

Fig.1  Seven examples of image categories

下载: 原图 | 高精图 | 低精图

2.2　数据增强

数据增强是一种使有限的数据产生更多的等价数据来人工扩展训练数据集的技术［23］。一般来说，卷积神经网络（CNN）的准确率随着训练图像的数量增多而增加。当难以获得足够的训练图像时，数据增强成为扩展数据库的有效手段，可以有效提高CNN的分类精度［24］。数据增强最主要的目的是提高训练数据的多样性，从而帮助模型更好地泛化到看不见的测试数据，并且可以很好地解决由于样本不足而出现的过度拟合问题。通过数据增强方法可以很好地对未来我国的火星车所采集的岩石数据进行泛化，无需参数调整就可直接迁移使用。常用的数据增强方法可分为空间几何变换、像素颜色变换和多样本合成几种方法。本文使用的图像增强技术有：翻转、模糊、裁剪、旋转、平移和加入噪声。本实验中数据增强前后数据对比如图2所示。

图2  数据增强前后的类别数量对比

Fig.2  Comparison of the number of categories before and after data enhancement

下载: 原图 | 高精图 | 低精图

3 模型压缩

传统深度神经网络模型计算量大且内存密集。然而，火星车探测车的系统内存资源有限，分类任务不能占据系统的全部资源。大计算量阻碍了深度神经网络模型在内存资源低的设备或具有严格延迟要求的应用程序中的部署。在不降低精度的前提下，压缩和加速深度网络模型，完成火星图像分类是必然的途径。

目前，通常意义上的模型压缩，主流思路有两个方向。一个是设计更轻巧的网络模型，如ShuffleNet、MobileNet和EfficientNet等。这些轻量级网络模型，可以根据实际精度要求，不断调试优化。另外一种思路是通过模型压缩方法改变网络模型结构，减小模型的内存占用量和计算复杂度。表2给出了4种模型压缩及其简要描述。模型压缩的思路是针对ResNet［25］、DenseNet［26］、VGGNet等网络模型中存在大量冗余节点，节点权重占用大量内存的问题，通过剪除冗余节点以达到优化的目的。

表2  模型压缩方法与简要介绍

Tab.2  Model compression method and brief introduction

方法 简要介绍

知识蒸馏 通过采用预先训练好的复杂模型（Teacher model）的输出作为监督信号去训练另外一个简单的网络（Student model），可以在卷积层和全连接层操作，仅支持从头开始训练。

低秩分解 使用矩阵/张量分解来估计信息参数。支持卷积层和全连接层操作，仅支持从头开始训练。

转移/紧凑卷积核 设计特殊的结构卷积核来保存参数。仅在卷积层操作并且只支持从头开始训练。

迭代剪枝与量化 减少对性能不敏感的冗余参数。可以在卷积层和全连接层操作，支持从头开始训练和预训练模型。

下载: 导出CSV

与其他压缩模型方法相比，如知识蒸馏、低秩分解和转移/紧凑卷积核，网络剪枝在保留精度的同时大幅减少了计算复杂度。因此，本文采用迭代剪枝的方法。

3.1　迭代剪枝

迭代剪枝方式可以降低卷积神经网络的计算量和内存占用量。VGG-16具有较好的图像分类效果［27］，但VGG-16参数量大，高达113.84 Mb，消耗了大量内存。VGG-16拥有13层的卷积层和3层的全连接层。VGG-16中参数量最多的层并非卷积层，而是第一个全连接层。这是因为13层卷积层提取的特征参数全部传输到第一个全连接层，共有84.24 Mb参数，占总参数量的74%。

降低全连接层的参数量的关键步骤是确定每层中最佳的神经元数量。然而，全连接层的神经元数量通常根据工程经验确定，如128、512、4 096等。通过控制剪枝率对模型中神经元数量进行部分修剪，少量修剪并不会造成分类精度的下降。这种被修剪的弱神经元是多余的，其存在会导致过拟合和优化困难。为了寻找每层神经元的最佳数量以实现更高效的网络架构，本文采用一种迭代剪枝方法以修剪冗余的神经元。

迭代剪枝分为3步，如图3所示。

图3  迭代剪枝的流程图

Fig.3  Flowchart of iterative pruning

下载: 原图 | 高精图 | 低精图

第一步，利用迁移学习训练VGG-16的连通性。与传统训练不同的是，训练VGG-16的连通性不是学习权重值，而是学习神经元的重要性。

第二步，修剪低权重连接。设训练数据集为D=｛X=｛x0，x1，…，xN｝，Y=｛y0，y1，…，yN｝｝，其中X、Y代表输入和输出。优化权重参数W= {(w11,b11),(w21,b21),...(wC𝓁L,bC𝓁L)}，使得代价函数值最小，如式（1）所示：

J(θ)=−1m⋅[∑i=1m∑k=1K(y(i)kloghθ(x(i))+(1−y(i)k)log(1−(hθ(x(i)))k)]

 . （1）

权重标准差Θstd(W)的计算公式如式（2）所示：

Θstd(W)=1|W|∑i(wi−μw)2−−−−−−−−−−−−−−−−√

 ， （2）

其中，W是权重集的维数。

将权重低于Θstd(W)的神经元h置0。此时代价函数为J(θ',h=0)。需使剪枝前后网络模型的代价函数尽可能相等，因此需给出使函数argmin|J(θ)−J(θ',h=0)|最小的h，从而将密集网络转为稀疏网络。

第三步，训练稀疏网络，以便剩余的连接可以补偿已删除的连接，这一步至关重要。修剪和再训练的阶段可以迭代，减少复杂度。否则，准确性会下降。

修剪连接后，零输入连接或零输出连接的神经元可以被安全地修剪，修剪示例如图4所示。在训练过程中，死亡的神经元将被自动移除。

图4  修剪前后的突触和神经元

Fig.4  Synapses and neurons before and after pruning

下载: 原图 | 高精图 | 低精图

3.2　量化与Huffman编码

剪枝是删除网络模型中不重要的权重参数，量化则是减少存储权重的字节数，将权重存储在更小的空间中。权重不需要以高精度方式存储，甚至不应该以高精度方式存储。本文采用K-means++聚类实现权重参数的量化，利用霍夫曼编码压缩迭代剪枝与量化后的VGGNet权重参数，达到减少存储量和浮点数运算量（Flops）的作用。在利用霍夫曼编码实现压缩时［28］，概率大的分配短码，概率小的分配长码。

4 实验结果与分析

本节将在Mars32K数据集上评估迭代剪枝与量化对VGG-16进行模型压缩的效果。

采用5种数据增强方法进行数据扩充，并且将图像的像素都调整为224×224。本文使用迁移学习方法训练VGG-16模型的连通性。在迭代剪枝过程中，使用SGD优化器以10-4的学习率、批量大小为32、动量为0.8、剪枝率为0.2训练网络模型。

本文采用准确率、浮点数运算量和总内存占用量指标来评价模型的性能。用TP、FP、FN、TN分别表示标签为正样本、预测为负样本，标签为正样本预测、为正样本，标签为负样本预测、为负样本和标签为负样本预测、为正样本，准确率的计算公式如式（3）所示：

Accuracy=TN+TPTN+TP+FN+FP

 . （3）

浮点数运算量的计算公式如式（4）所示：

Flops=2HW(CinK2+1)Cout

 ， （4）

式中：H、W和Cin分别表示输入特征图的高度、宽度和通道数，K是卷积核的宽度，Cout是输入特征图的通道数。

4.1　实验结果

本文探索了不同剪枝率对模型准确率的影响。从图5可以看出，随着剪枝率的增加，分类准确率下降。值得注意的是，图5中修剪后不再训练，剪枝后的网络性能未达到最优。鉴于剪枝率为0.2时，VGG-16效果较好，本文选择0.2的剪枝率，并在后续实验中，剪枝后再训练网络。

图5  不同剪枝率对模型准确率的影响

Fig.5  Influence of different pruning rates on model accuracy

下载: 原图 | 高精图 | 低精图

模型压缩分为3步。第一步，迁移学习预训练VGG-16的连通性，结果对应图6中前150代的预训练模型曲线；第二步，迭代修剪低权重神经元，结果对应图6中迭代剪枝曲线；第三步，对修剪后的网络再训练，结果对应图6中再训练曲线。

图6  预训练、迭代剪枝和再训练模型的准确率。

Fig.6  Accuracy of pretrained， iteratively pruned and retrained model.

下载: 原图 | 高精图 | 低精图

为了验证迭代剪枝再训练VGG-16的有效性，将训练代数为160、200和275时的权重参数进行量化与霍夫曼编码。图7给出了训练代数为160时的结果。图8给出了训练代数为160、200和275时对应的特征提取图。表3给出了模型压缩前后的相关数据。从图8和表3可以看出，一方面，随着训练代数增大，未裁剪的神经元数量变少，提取的岩石相对位置信息和纹理信息减少，分类准确率也随之下降；另一方面，内存占用量、Flops值也随之降低。对以上指标折中，本文选择训练代数为160的压缩后的VGG-16。

图7  霍夫曼编码量化实现权重参数压缩的过程

Fig.7  Process of Huffman coding and quantization to realize the compression of weight parameters

下载: 原图 | 高精图 | 低精图

图8  迭代剪枝模型不同层提取的特征图

Fig.8  Feature maps extracted from different layers of iterative pruning model

下载: 原图 | 高精图 | 低精图

表3  模型压缩前后的数据对比

Tab.3  Data comparison before and after model compression

方法 训练代数 内存占用量/Mb Flops/MFlops 准确度/%

VGG-16（预训练模型） 150 241.99 346.20 98.10

迭代剪枝 160 60.46 148.60 92.82

200 25.43 63.80 85.54

275 17.48 36.45 62.79

再训练 160 62.63 150.60 96.15

200 29.61 69.70 91.54

275 23.67 43.21 67.37

下载: 导出CSV

通过混淆矩阵展示模型压缩后的性能，如图9所示。通过分析混淆矩阵可以看出，分类准确率高，仅存在极少失败案例。如迭代剪枝模型将G6类别错误地分类为G2类，G2类别被错误地分类为G4类别。

图9  分类混淆矩阵

Fig.9  Classification confusion matrix

下载: 原图 | 高精图 | 低精图

综合考虑模型的内存占用量、准确率和Flops 3个指标，以160代时再训练的权重参数进行量化与霍夫曼编码。编码后的权重参数构成的网络作为火星图像分类模型。

4.2　消融实验

为了验证迭代剪枝的有效性，本文将VGG-16与轻量级神经网络做对比，如MobileNet、ShuffleNet、EfficientNet、高分类精度的残差神经网络ResNet50和采用密集连接的DenseNet。

MobileNet核心思想是采用深度卷积操作。相较于标准卷积操作，可以减少数倍的计算量，从而达到提升网络运算速度的目的。然而采用深度卷积会出现“信息流通不畅”问题，即输出的特征图仅包含输入的特征图的一部分。

针对深度卷积带来的“信息流通不畅”问题，ShuffleNet采用组卷积和通道随机混合操作两种经典方法，其准确率得到提升。

EfficientNet并非真正意义上的轻量化模型，只是借鉴depth-wise convolution，而depth-wise convolution又是上述二个轻量化模型的关键点。MobileNet、ShuffleNet是真正意义上的轻量化网络，而EfficientNet在同等参数数量条件下可获得更高的性能。

从表3和表4可以看出，在这5种网络中，EfficientNet内存占用量最低，ShuffleNet的Flops最小，均大于压缩后的VGG-16。ResNet50模型准确率最高，小于压缩后的VGG-16。因此，压缩后的VGG-16模型更轻量，准确率较高，适合部署到未来的火星探测任务中。

表4  消融实验

Tab.4  Ablation experiment

算法 内存占用量/Mb Flops/MFlops 准确度/%

MobileNet 162.20 320.24 98.71

ShuffleNet 139.60 316.45 98.43

EfficientNet 123.40 399.30 98.12

ResNet50 317.43 4 218.88 98.86

DenseNet 147.10 2 880.00 98.30

下载: 导出CSV

5 结论

本文设计了一种轻量级的图像分类器。首先，采用迁移学习训练网络的连通性。其次，通过迭代剪枝方法修剪不重要的神经元。最后，采用K-means++聚类实现权重参数的量化，利用霍夫曼编码压缩VGGNet权重参数。此外，采用5种数据增强方法进行数据扩充，提高训练数据的多样性。压缩后的VGG-16具有如下优点：

（1）准确率高。压缩后的VGG-16模型分类精度为96.16%。这是因为迭代剪枝未损害模型的特征提取能力，仅裁剪冗余神经元。剪枝后的神经元能够完整地提取火星岩石的相对位置信息和纹理信息。

（2）内存占用量和浮点数运算量小。压缩后的VGG-16模型内存占用量和Flops值分别是62.63 Mb和150.60 MFlops，均低于5种对比网络。

（3）泛化能力和鲁棒性强。压缩后的VGG-16可在参数未调整的情况下对岩石进行分类，可较好地应用于未来的火星探测任务。