深度学习轻量化侦察图像压缩网络

谌 宇,谌德荣,崇魁奇,王泽鹏,张 凯

(1.北京理工大学,北京 100081;2.北京航宇天穹科技有限公司,北京 100043)

0 引言

小型侦察平台具有灵活性高、隐蔽性强等特点,是快速获取侦察图像的重要技术手段。为了保证获取图像的时效性,小型侦察平台需要将获取的侦察图像压缩后通过无线信道传回指挥中心,图像压缩性能是决定获取图像质量的关键。

根据处理方法的不同,图像压缩算法可以分为传统压缩算法和基于深度学习的压缩算法两类。

传统压缩算法主要通过预测、变换去除图像的空间冗余,通过熵编码去除统计冗余,进而达到压缩数据的目的。例如,JPEG2000[1]压缩算法通过小波变换来去除图像中的空间冗余,利用算数编码来去除统计冗余,是传统图像压缩算法中性能最好的。

基于深度学习的图像压缩算法主要通过神经网络来去除图像中的冗余信息。根据使用的网络模型结构不同,基于深度学习的图像压缩方法可以分为基于递归神经网络的、基于变分自编码器的和基于生成对抗网络的图像压缩算法三类。

基于递归神经网络(RNN)的图像压缩算法通过迭代调用基于RNN的编码器来压缩图像或图像残差的数据量,并利用熵编码进一步降低数据量。虽然能精确地控制压缩比,但是存在网络模型复杂度较高、计算量大、图像压缩质量不佳等问题。

基于变分自编码器(VAE)的图像压缩算法利用卷积神经网络搭建由编码器和解码器组成的端到端学习网络。此算法可以超过JPEG2000的图像压缩质量,但是单个网络模型只能实现一种压缩比的图像压缩,而且使用了熵编码处理,不利于硬件的并行加速处理。

基于生成对抗网络(GAN)的图像压缩算法在基于VAE的端到端图像压缩模型的基础上引入判别器以构成GAN,并且以金字塔式结构设计了编码器和判别器,在去除图像冗余的同时保留更多细节信息,提高了图像质量,但同样存在模型复杂、计算量大等问题。

为满足小型侦察平台对图像压缩算法计算复杂度低、计算量小、模型参数少等要求,提出基于深度学习轻量化侦察图像压缩网络,通过轻量化编码网络设计、高质量解码网络设计和损失函数设计以降低编码网络的复杂度和提升图像压缩质量。该网络在编码端采用三个卷积模块将图像直接映射为服从均匀分布的二进制码流,得到压缩后数据,进而去除了不利于硬件并行处理的熵编码操作;为了降低卷积模块的计算量和参数量,采用深度可分离卷积代替传统卷积,并利用分组卷积+通道重排优化卷积模块;在解码端通过转置卷积和残差模块提升特征的提取和恢复能力,进而提高解码图像质量。为了使压缩后的数据服从均匀分布使其包含的信息量最大化,并使图像重建误差最小化,设计了能够同时反映压缩后数据信息熵和重建误差的目标函数。

1 卷积神经网络

卷积神经网络(CNN)是一种包含卷积计算且具有深层次结构的神经网络,通常由数据输入层、卷积层、激活函数层、池化层与全连接层组成。

数据输入层一般是对输入卷积网络的数据进行预处理,例如去均值、归一化等。该层的主要目的是加快训练过程中网络的收敛速度,调整输入数据的取值范围,防止超过激活函数的值域限制。

卷积层主要通过卷积核完成对数据的卷积操作。该层的目的是获取数据更深层次的特征信息。在卷积神经网络中,决定某一层输出结果中一个元素所对应的输入层的区域大小,被称作感受野。卷积层的层数越多,感受野就越大;感受野越大,网络能提取的信息就越充分。

激活函数的主要作用是向网络中引入非线性,加入激活函数之后,深度神经网络具备了非线性映射学习能力。

全连接层一般处在网络的末端,在整个卷积神经网络中起到“分类器”的作用,将网络中生成的各种特征整合到一起。 该层的参数较多,容易出现参数冗余,降低训练速度等问题。

2 轻量化图像压缩网络

通过轻量化编码网络设计、高质量解码网络设计和损失函数设计以降低编码网络的参数量与计算量和提升图像压缩的质量。

2.1 编码网络设计

编码网络的设计以编码结构简单、编码参数少、编码计算量少为指导思想。编码网络使用卷积神经网络将图像直接映射为二进制码流,同时为了减少编码的计算量,引入深度可分离卷积与分组卷积,大大降低了网络的参数量与计算量。编码网络结构如图1所示。

图1 编码网络结构示意图Fig.1 Schematic diagram of coding network structure

编码网络主要包含三个模块:

第一个模块包含一个下采样卷积层和一个ReLU激活函数。下采样卷积层对输入数据进行步长为2的下采样卷积处理。该层卷积核尺寸为3×3,卷积核个数为32个。

第二个模块采用了深度可分离卷积以及分组卷积+通道重排。该模块包含一个ReLU激活函数、一个下采样深度卷积层和一个点卷积层。下采样深度卷积层,进行步长为2的下采样,其卷积核为尺寸3×3,个数为32个。点卷积层根据分组卷积原理将该层分为两个部分,每部分卷积核尺寸为1×1×16,卷积核个数为16个。

第三个模块也采用了深度可分离卷积。该模块也包含一个下采样深度卷积层,以及一个点卷积层。下采样深度卷积层采样步长为2,卷积核尺寸为3×3,卷积核个数为32个。点卷积层卷积核尺寸为1×1×32,该层卷积核数量与压缩比有关。

经过三个模块处理后,输出数据经二进制量化得到压缩后的数据。二进制量化处理将网络输出的数据量化为位宽为1 b的0或1,进一步提高压缩比。在测试过程中二值量化采用以0.5为分界的阶跃函数;在训练过程中为了保证梯度能够反传,采用Sigmoid函数代替阶跃函数完成量化。

压缩比与第三模块点卷积层卷积核的数量相关。设编码网络输入的灰度图像长和宽分别为H,W,第三模块点卷积层的卷积核个数为M,则编码网络输入的数据量为H×W×1,编码网络输出数据量为H/8×W/8×M,此时数据量压缩比为64/M。原始图像像素位宽为8 b,二值量化后数据位宽为1 b,因此最终压缩比为64/M×8。当输入图像尺寸为128×128,压缩比设定为64时,第三模块点卷积层个数为8,编码端参数量为1 632,编码器浮点运算量(FLOPs)为4 069 376(一次乘加运算FLOPs为2)。

2.1.1深度可分离卷积

在编码网络中,第一模块输入通道数为1,第二、第三模块输入通道数为32。计算量主要集中在第二模块和第三模块。在这两个模块中使用深度可分离卷积,可以大大减少编码网络的参数量与计算量[2-3]。

深度可分离卷积包括深度卷积(DW)和点卷积(PW)两个部分。相比常规的卷积操作,其参数量和计算量比较低。

在深度卷积中,每个输入通道的特征图分别与一个卷积核进行卷积,而非每个输入通道的特征图与所有卷积核进行卷积。深度卷积的计算结果缺少不同通道在相同空间上的信息,通常需要在其后加一个点卷积,以融合不同通道之间的信息。深度卷积的卷积核数与输入通道数一致,不会改变数据的维度。

点卷积的卷积核尺寸为1×1×M,M为输入通道数。点卷积运算主要是对每个输入特征图的像素点在通道的维度上进行线性组合。

若编码网络第二、第三模块不使用深度可分离卷积,则编码网络结构如图2所示。

图2 剥离深度可分离卷积的编码网络结构Fig.2 Schematic diagram of code network structure without using depth-separable convolution

图2中,第二模块只保留了一层下采样卷积层和一个ReLU激活函数。卷积层分为两个部分,每个部分卷积核尺寸为3×3×16,卷积核个数为16个;第三个模块也只保留了一层下采样卷积层,该层卷积核尺寸为3×3×32,卷积核个数为8个。此时编码网络参数量为7 200,是使用深度可分离卷积的编码网络参数量的4.4倍;在输入图像尺寸为128×128的情况下,浮点运算量为11 750 400,是使用深度可分离卷积的编码网络运算量的2.9倍。

2.1.2分组卷积+通道重排

在编码网络三个模块中第二模块的输入通道与卷积核的数量最多,可通过对该模块采用分组卷积[4-5]+通道重排[6-8]的方法以减少参数量及计算量。

分组卷积是将输入通道以及卷积核分为T组,组内进行常规卷积,组与组之间没有联系。其中,每一组输入通道和卷积核个数为原来的1/T。

分组卷积在降低计算量的同时削弱了通道之间的相关性,引入ShuffleNet算法的通道重排模块可提高对通道相关性利用。分组卷积与通道重排结构如图3所示。

图3 分组卷积与通道重排示意图Fig.3 Schematic diagram of group convolution and channel shuffle

在编码网络的第二模块中,点卷积层使用了分组卷积,将该层卷积分两组,每组由16个1×1×16卷积构成,输入通道数为第二模块深度卷积输出通道数的一半,即每一组的输入通道数为16。

为了提高对通道相关性的利用,参考ShuffleNet算法,首先将第一模块下采样卷积层输出的32个通道进行分组,通道1～16为一组,通道17～32为一组;然后将两组通道进行交错排列,并将重排后的特征图传入下一个处理网络。

若编码网络第二模块不使用分组卷积+通道重排,则编码网络结构如图4所示。

图4 剥离分组卷积+通道重排的编码网络结构图Fig.4 Schematic diagram of code network structure without using Group convolution and channel shuffle

与图1相比,图4中第一、第三模块未发生变化,第二模块的点卷积层卷积尺寸变为了1×1×32,卷积核个数变为32个。此时编码网络参数量为2 144,是使用分组卷积+通道重排编码网络参数量的1.3倍;在输入图像尺寸为128×128的情况下,编码网络浮点运算量为5 117 952,是使用分组卷积+通道重排编码网络参数量的1.25倍。

表1、表2比较了在128×128分辨率下本文设计的编码网络、不使用深度可分离卷积的编码网络、不使用卷积分组+通道重排的编码网络,三种编码网络的参数量与计算量。参数量见表1,计算量见表2。

表1 编码网络参数量对比表Tab.1 Code network parameter number comparison table

表2 编码网络计算量对比表Tab.2 Comparison table of computations at coding network

由表1、表2可知,深度可分离卷积和卷积分组+通道重排可以使编码网络参数量、计算量大大降低。

2.2 解码网络设计

解码网络以提高重建图像质量为设计目的。提出的解码网络结构如图5所示。

图5 解码网络结构示意图Fig.5 Decoding network structure diagram

在图5中,解码网络主要由3种模块组成:卷积层、转置卷积、残差模块。

在解码网络中,第一个卷积层尺寸为3×3,卷积核个数为64个;最后一个卷积层尺寸为3×3,卷积核个数为1个;其余的卷积层卷积尺寸为3×3,卷积核个数为32个。

转置卷积用来恢复图像分辨率。转置卷积的核长为4,步长为2,卷积核个数为32个。这样可以避免由棋盘效应带来的重建图像质量下降的问题。

残差模块可以缓解网络过深造成的网络退化和梯度消失问题,有利于提升梯度传递的效率。残差模块结构中包含两层卷积层和两个ReLU函数。

2.3 损失函数设计

(1)

式中:压缩量y的熵值如式(2)所示。

(2)

式中:n为压缩量y的取值范围,pi为对应y值出现的概率。

(3)

3 网络训练与测试

网络训练所用平台配置如下:CPU为Intel Xeon(R)Silver 4210R @2.40 GHz×40,GPU为 NVIDIA RTX 3090 TI,内存为64 GB。

3.1 数据集准备

实验数据来自2018和2019届CLIC图像压缩挑战赛,共616张图像。其中随机选择587张原始图像随机裁剪为尺寸为256×256的图像,共获得24 078张图像用于构成训练集。剩余42张原始图像,裁剪为6 692张大小为128×128的图像作为验证集。选用Set5,Set14[9],SunHays80,BSD100,Urban100[10]作为本次实验的测试集,将测试集图像裁剪为128×128的图像进行测试。

3.2 模型训练

使用Adam优化器进行训练,初始学习率设置为0.000 1,在达到第80轮的时候学习率降低10倍,共训练120个轮。模型效果的评估准则为PSNR和SSIM[11],PSNR计算公式为

(4)

式中:(I,K)MSE表示图像I和K残差值的平方。

(5)

SSIM计算公式为

(I,K)SSIM=L(I,K)·C(I,K)·S(I,K),

(6)

(7)

(8)

(9)

式中,uI和uK分别表示图像I和K的均值;σI和σK分别表示图像I和K的标准差;σIK表示图像I和K的协方差;C1,C2,C3为常数,避免分母为0。

3.3 测试结果

对于测试集,轻量化侦察图像压缩网络与JPGE2000算法编码所需时间如表3所示。

表3 编码时间对比表Tab.3 Code time comparison table

由表3可知,本文方法编码速度约是JPGE2000算法编码速度的11倍。

JPGE2000算法和本文算法重建图像质量如表4所示。选用5个测试集PSNR与SSIM的加权平均作为重建质量的度量。

表4 重建质量对比Tab.4 Reconstruction quality comparison

由表4可知,轻量化侦察图像压缩网络与JPGE2000算法相比,PSNR提高了3.85 dB,SSIM提高了0.27,重建质量远超于JPGE2000算法。解码端使用转置卷积能获得更好的重建图像质量。

图6给出了JPGE2000和本文算法重建侦察图像(分辨率为128×128)的结果。进一步说明了轻量化侦察图像压缩网络的重建质量高于JPGE2000。

图6 128×128分辨率下图像对比Fig.6 Image comparison at 128×128 resolution

为了验证深度可分离卷积、分组卷积+通道重排以及转置卷积的作用,本文进行了剥离实验以验证其效果。对比了用普通卷积替换深度可分离卷积、普通卷积替换分组卷积、最近邻插值替换转置卷积三种情况与本文方法的重建图像精度,结果如表5所示。

表5 重建质量对比Tab.5 Reconstruction quality comparison

对比表1、表2以及表5中无深度可分离卷积、无分组重排以及轻量化侦察图像压缩网络可知,引入深度可分离卷积与分组卷积+通道重排不仅可以减少编码网络的参数量与计算量,还能提高重建图像的质量。用转置卷积进行上采样比用最近邻插值进行上采样能获得更好的重建图像质量。

4 结论

本文设计了一种基于深度学习轻量化侦察图像压缩网络,编码网络由三个卷积模块组成,并通过深度可分离卷积、分组卷积+通道重排的方式降低编码网络的参数量与计算量。该编码网络能够将输入数据直接映射为服从均匀分布的二进制码流,从而完成图像压缩;解码端使用转置卷积、残差模块增加特征的提取与重建能力。实验结果表明,与JPGE2000算法相比,基于深度学习轻量化侦察图像压缩网络的PSNR提高了3.85 dB,SSIM提高了0.27,编码时间降低91%。