基于YOLO-v3模型压缩的卫星图像船只实时检测

陈科峻， 张 叶

( 1. 中国科学院 长春光学精密机械与物理研究所 应用光学国家重点实验室，吉林 长春 130033；2. 中国科学院大学， 北京 100049)

1 引 言

随着海上运输业的快速发展，海上遇难船只搜救、船舶避撞、非法捕鱼、非法货运船只的管控等任务，对海上船只精确定位与识别技术有着广泛的需求和重要的现实意义[1]。因此，从环境保护机构到国家政府当局，许多组织对海上船只进行更密切的监视。随着卫星技术的日益成熟和卫星航空成像精度的日益提高，卫星图像中包含的信息多样，利用卫星图像进行海上船只的检测受到了广泛的关注。然而，视点、旋转、背景等多种因素变化的影响使得卫星图像船只检测任务依然存在严峻的挑战。

目前，国内外研究人员对图像船舰目标检测做了大量的相关研究。文献[2]通过分析遥感图像中的明暗像素的关联参数对船只进行目标检测，实现图像中明暗场景的快速识别与分离。然而基于传统特征提取的方法进行目标检测往往精度不高。由于海洋表面微波散射情况复杂，传统船舰检测算法能力有限。文献[3]针对以上问题，通过分辨率归一化制作混合的训练样本数据集，基于Faster-RCNN[4]的框架下，构建了一个3层卷积神经网络的船舰检测算法，取得了较好的检测效果。但基于Faster-RCNN的船舰检测算法模型复杂度较高，检测速度难以达到实时性检测的要求。近年来，随着基于深度学习的目标检测技术的发展，YOLO(you only look once)[5]作为一个兼顾了检测速度和检测精度的算法被广泛应用于日常的目标检测任务中。文献[6]提出基于YOLO 的车流量检测方法，可以较好地对混合车流中的车辆进行跟踪和统计。但基于原始YOLO算法的目标检测模型参数量较大，对嵌入式设备计算性能要求较高。文献[7]提出了一种基于改进 YOLO模型的舰船目标识别方法。通过精简原始模型的卷积层，减少模型的参数量。同时在训练过程中引入迁移学习的概念，防止模型过拟合，达到减少计算资源开销的目的。但该方法主要针对大型船舰的目标检测，对于卫星图像中的小目标船舰检测效果不理想。基于端到端回归的YOLO-v3[8]网络在YOLO算法的基础上进行改进，在保证快速检测的同时，也达到了较高的检测精度，使得基于深度学习的目标检测算法的落地应用成为可能。

本文将YOLO-v3目标检测网络应用于卫星图像船舰目标检测任务中。针对输入图像中船舰的外形特点，构建图像金字塔对训练输入图像进行多尺度变换，丰富原始图像的特征。由于YOLO-v3的初始锚点框[9](anchor)的尺寸是针对VOC[10]等数据集计算设计得到的，并不适用于文本的卫星船舰图像数据集。本文采用K-means(K均值)聚类算法的改进版K-means++算法进行初始锚点框数量和尺寸大小的聚类，构建出适合船舰数据集的船舰目标检测网络结构。最后，为了能够将模型应用于嵌入式设备中，对训练好的模型进行模型剪枝，将模型中冗余的通道进行裁剪，减少模型的参数量，提高了模型的检测速度。

2 图像船舰检测原理

2.1 多尺度金字塔图像

针对船舰的图像特征，本文采用多尺度金字塔图像作为模型训练的输入。本文受到文献[11]的启发，采用高斯金字塔与拉普拉斯金字塔相结合的方式，先采用高斯金字塔对图像进行下采样以获得新的图像，用拉普拉斯金字塔对图像进行上采样重建获得新样本。

2.1.1 高斯金字塔

对于大小为w×h的图像I，高斯金字塔Gi由I的几个分辨率减小的高斯图像Ii组成。Ii的尺度计算公式如下：

(1)

式中，i= (0,1,...,j)代表金字塔的级数。

图像Gi通过高斯函数生成的卷积核进行卷积。公式如下：

(2)

式中，m,n代表原图像的尺度，L代表金字塔的层级，w代表窗口函数。

2.1.2 拉普拉斯金字塔

拉普拉斯金字塔由高斯金字塔计算得来，公式如下：

Li=Gi-up(Gi)⊗g5×5，

(3)

式中，Gi表示第i层的图像，而up操作表示上采样。 ⊗表示卷积运算，g5×5表示5×5的高斯卷积核。最终得到的图像金字塔效果图如图1所示。

图1 多尺度图像金字塔Fig.1 Multi-scale image pyramid

2.2 YOLO-v3目标检测原理

YOLO-v3算法将目标区域检测与物体类别预测看作回归问题。YOLO-v3算法的特征提取网络为Darknet53。Darknet53包含52个卷积层，用于对图像进行特征提取，为了防止深层网络训练过程中发生模型退化，YOLO-v3借鉴了ResNet[12]的做法，在层与层之间进行了链路的连接。最后一层为1个全连接层。YOLO-v3通过单个神经网络，即可实现对图像中目标的定位和分类。Darknet-53网络结构如图2所示。

图2 Darknet53网络结构

Darknet-53网络对图像进行了5次下采样，每次采样的步长为2。其中卷积模块的结构如图3。

图3 Darknet53网络结构Fig.3 Module of Darknet53 convolution set

YOLO-v3算法假设目标散落在图像中各个区域，因此将图像划分为多个子区域。接着对各个区域预测多个锚点框( Bounding Box )，并求出每个边框的置信度(Confidence)。置信度分数的大小，反映了锚点框内存在目标物体的准确程度。置信度定义如下式：

(4)

式中，Pr ( Object )表示锚点框中存在目标图像的概率大小。IoU表示预测的目标框与真实的标注框之间的交并比，IoU的计算公式如下：

(5)

式中，C代表预测出物体的边界框，G代表实际标注的物体边界框(Ground Truth)。

2.3 初始锚点框聚类算法

常见的目标检测算法中，需要预先对初始的锚点框(Anchor)进行设定。锚点框与检测目标的契合度对目标检测算法的精度产生极大影响。针对卫星图像船舰的形状特征，本文采用聚类算法对YOLO-v3的锚点框进行初始化设定。

理想的聚类距离度量标准应该是计算误差与锚点框尺寸大小没有太大的关系，因此定义IoU距离作为聚类算法距离的度量。IoU距离如式(6)所示：

dIoU(box,centroid)=1-IoU(box,centroid)，

(6)

式中， IoU( box, centroid )表示某一锚点框box与聚类中心锚点框的IoU，dIoU即代表计算得到的IoU距离。

本文采用K-means++聚类方法对YOLO-v3的初始锚点框进行聚类。传统的K-means 算法通过随机初始化的方式选定初始聚类中心，往往容易陷入局部最优，得到的聚类结果无法反映真实的数据分布情况。K-means++聚类算法在K-means算法基础上进行了改进。具体计算步骤如下：

(1)在数据集中随机初始化选择1个点作为第一个簇的中心点；

(2)计算数据集中各个样本点到每个已知的簇中心的距离D(x)；

(3)定义概率Px表示每个样本被选为下一个簇的中心的概率：

(7)

(4)重复(2),(3)步骤，直至确定出K个聚类中心；

(5)将第(4)步得到的K个聚类中心作为初始聚类中心点，代入传统的K-means算法中，得到最终的聚类中心点。

从K-means++算法的计算过程可以看出，聚类中心的初始化以距离作为依据，和当前聚类中心距离较大的点，属于当前类的概率则比较小，因此有更大的可能性被选为新的聚类中心，聚类中心的初始化更为合理。

3 基于YOLO-v3的模型剪枝

图像船舰检测任务的模型往往被部署在计算能力有限的设备上，对模型的检测时间和模型的参数量大小有着严格的要求。因而本文对原始的模型进行模型剪枝压缩。

模型剪枝是模型压缩的常用方法，对深度神经网络的稠密连接引入稀疏性，通过将“不重要”的权值直接置零来减少非零权值数量。典型的模型剪枝方法包含4个步骤[13]：(1)分析预训练模型中每一个神经元的重要程度；(2)剔除模型中对模型推理激活程度较低的神经元；(3)微调模型，提高修剪后模型的精度；(4)对微调后的模型进行测试，确定修剪后的模型是否满足要求。模型剪枝的流程如图4所示。

图4 模型剪枝流程图Fig.4 Flow chart of model pruning

3.1 稀疏训练

对YOLO-v3模型进行压缩，需要对训练好的模型进行稀疏化训练。稀疏训练的目的即在模型训练过程鉴别重要性较低的通道，以便于对重要性较低的通道进行裁剪。由于YOLO-v3网络的每个卷积层后面都有一个批归一化层来加速网络收敛和改进泛化性能。本文将批归一化层(Batch Normalization)[14]中的尺度因子作为通道重要性的指标，并对这些尺度因子计算L1范数并融入损失函数中训练，得到每个通道的重要性分数[15]。应用于精简基于卷积神经网络(CNN)的图像分类器，可以显著地减小模型参数量。

批归一化层的更新过程如下。求解上一层输出数据的均值和标准差：

(8)

(9)

(10)

式中，ε是为了避免分母为0而加入的一个趋近于0的值。经过对上面归一化处理得到的数据进行重构即可得到：

(11)

式中，γ和β是可以学习的参数，用于将归一化后的数据分布进行还原。

本文将γ值作为通道重要性的衡量标准，对每一个通道的γ值计算L1范数，具体计算公式如下：

(12)

最终模型训练的损失函数在YOLO-v3的损失函数基础上增加了对γ的正则化项：

(13)

式中，lossyolo表示YOLO-v3的损失函数，α表示惩罚项。

3.2 模型剪枝

4 实验与结果分析

4.1 数据集准备

基于深度学习的目标检测往往需要大量的数据用于模型的训练。为了检验本文方法的有效性，本文采集到网络上的船舰图像数据集[16]，该数据集通过卫星进行图像采集。卫星船舰数据集分辨率大小为768×768像素，对图像进行了标注。最终得到精确标注的5 000张卫星船舰图像用于模型的训练。部分图像样例如图5所示。

图5 卫星船舰图像样例Fig.5 Example of satellite ship images

4.2 锚点框聚类

本文对标注好的数据样本进行聚类，预先选好的锚点框聚类个数为k∈(1,…,10)。为了对比不同的聚类算法对模型的影响，在保证各训练参数固定的情况下，分别采用K-means和K-means++对锚点框进行聚类。表1为不同聚类算法下的部分锚点框聚类尺寸。

表1 初始锚点框聚类结果Tab.1 Results of initial candidate box clustering

为了评价不同聚类算法的性能，采用平均准确率(mAP)作为评价指标，从整体上反映了目标定位算法的精度。平均准确率综合了检测模型的精确率(Precious)和召回率(Recall)：

(14)

(15)

(16)

式中：TP表示真正例，FP表示假正例，FN表示假负例，Pre表示精确率，Rec表示召回率。不同个数的锚点框下验证集的平均准确率如图6所示。

图6 不同聚类算法结果比较Fig.6 Results of different clustering algorithms

从图6可以看出，采用K-means++进行锚点框度量可以提高目标检测算法的定位精度。所得到的最高平均准确率高于采用K-means进行聚类所得到的最高平均准确率。当锚点框个数为6时，所得到的验证集的平均定位精度最高，为78.54%。相比于原始的YOLO-v3采用的9个锚点框，更少的锚点框可以提高模型的训练速度，减少模型的推断时间。

4.3 模型训练

对模型进行训练的实验系统环境为ubuntu16.04，基于Pytorch深度学习框架，显卡为NVIDIA 1660 Ti。本文对模型进行了2 000次迭代训练，初始学习率为0.001，采用动量梯度下降法，设置momentun为0.9。将原始YOLO-v3的初始锚点框替换成经过本文聚类之后得到的锚点框。模型训练过程损失曲线如图7所示。

图7 基础训练损失变化曲线Fig.7 Loss curve of basic training

从图7可以看出，经过2 000步的迭代训练，模型已经趋近于收敛，最终的损失稳定在0.665。

4.4 模型压缩

模型剪枝之前需要对模型进行稀疏训练，将批归一化层中的尺度因子作为衡量通道重要性的指标。实验中将稀疏训练的惩罚因子α分别设置为0.01，0.001，0.000 1。稀疏训练后的尺度因子分布直方图如图8所示。

原始模型的批归一化层因子γ近似于均值为1的正态分布。从尺度因子分布直方图可以看出，对模型稀疏训练的结果可以使得通道重要性较低的通道γ值压缩至接近于0，因此可以将γ值接近于0的通道进行裁剪。本文采用不同的惩罚因子对模型进行300个Epoch的稀疏训练，稀疏训练的损失曲线如图9所示。

(a)α=0.01

(b)α=0.001

图9 不同惩罚因子稀疏训练损失曲线Fig.9 Loss curve of sparsity training with different penalty factor

表2 剪枝前后模型参数对比

4.5 实验结果分析

通过对批归一化层中尺度因子大小接近于0的通道进行裁剪，剪枝后的模型参数量下降为原始模型的0.15倍，模型检测的平均准确率下降了4.13%。为了提高剪枝后模型的精度，采用微调的方式，对剪枝后的模型进行100个Epochs的微调训练，最终模型的平均检测精度稳定在77.31%，接近原始模型的精度。

为了验证本文卫星图像船舰检测方法的有效性，本文设置了对比实验，分别采用YOLO-v3，Faster R-CNN，SSD，R-FCN等常见目标检测算法对卫星图像船舰数据集进行检测。在同一台服务器上对以上各算法进行测试，测试输入图片大小为768×768像素。表3为各个算法的检测耗时与平均检测精度对比。

表3 各算法检测结果对比

从实验结果可以看出，本文模型的mAP达到了77.31%，在与其他检测模型检测精度相当的情况下，检测速度缩短到12 ms/张，模型参数量大小仅为原始YOLO-v3模型参数量的8.5%，对计算资源消耗更少，达到了实时性的目标检测要求。

5 结 论

本文针对卫星图像船舰检测问题，通过对训练图像进行多尺度金字塔变换，达到对模型进行多尺度训练的目的，提高了模型的特征提取能力。同时，采用K-means++聚类算法选定初始锚点框的个数和尺寸，提高了YOLO-v3模型的定位精度。接着对YOLO-v3目标检测模型进行模型剪枝。实验结果表明，本文的方法在保证检测精度的同时，减小了模型的参数量，检测时间由原来的30 ms/张缩短至12 ms/张，模型大小由原来的238.9 M缩小至21.5 M。满足了卫星、无人机等移动小算力设备的实时目标检测需求，为嵌入式目标检测算法的应用提供了参考依据。