自注意力SRGAN在岩石CT图像超分辨中的应用研究

朱联祥，郑 逸

(西安石油大学 计算机学院，陕西 西安 710065)

引 言

岩石显微图像可以在不改变岩石形态及内部结构的条件下，用于岩石的物性特征分析与计算，故而被广泛应用于石油地质领域[1]。然而在实际中，受限于采集设备成本及自然环境因素，一般难以获得足够高分辨率的岩石图像，从而对岩石分析的质量带来影响。单图像超分辨率技术可通过深度学习等方式，从单幅低分辨率图像构建出高分辨率的图像，具有节约成本且处理时间短的优点。

单张图像超分辨率重建算法主要经历了基于插值的算法、基于重建的算法和基于深度学习的算法3个阶段[2]。基于插值的算法是根据已知像素通过某种算法估计当前像素值，基于重建的算法则是在空域或频率域进行操作。针对这两种方法均存在的受图像内容限制及重建图像质量差的问题，Dong等人在2016年通过将深度学习引入到图像超分辨率研究，提出了SRCNN[3]。虽然所提算法只是采用了3层的卷积神经网络结构，但与双三次插值(Bicubic)及基于重建的算法相比，无论是在视觉效果，还是图像的峰值信噪比(PSNR)及结构相似度(SSIM)等方面，都有了明显进步。在此之后，亚像素卷积[4]、残差网络[5]、拉普拉斯金字塔[6]等网络结构陆续被应用于图像的超分辨率处理，并取得了较为显著的效果。

2017年，Ledig等人将生成对抗网络(Generative Adversarial Network，GAN[7])应用超分辨率领域，提出了SRGAN[8]。该方法将低分辨率的图像作为GAN的生成网络的输入，通过训练得到超分辨率的生成图像，并将其和原始高分辨率图像一起作为判别网络的输入。判别网络试图区分生成的图像和原始的高分辨率图像，而生成网络则试图通过训练网络参数来生成可以“欺骗”判别网络的超分辨率图像。通过判别网络和生成网络反复地迭代学习，最终达到纳什均衡，即生成网络可以生成足以乱真的超分辨图像，判别网络只能以50%的概率区分生成图像和原始高分辨率图像，从而实现低分辨率图像的超分辨处理。

最早提出并应用注意力机制的是自然语言处理任务[9]。随着时间的推移，注意力机制逐渐得到广泛应用。相较于传统的卷积方法，注意力机制可以通过得到的稀疏矩阵从全局范围内筛选出重要的部分区域，从而能更好地处理图像中的长范围依赖关系，协调生成图像的纹理细节，在一定程度上弥补卷积的不足。

本文以SRGAN为基础，通过引入自注意力机制[10]，对生成器和判别器的网络结构进行调整，以更好地表达高频特征。同时考虑到岩石显微图像的特殊性，通过在损失函数中引入岩石图像所特有的孔隙度特征来构建新的损失函数，实现了适合于岩石显微图像分辨率提升的超分辨率处理算法。最后用所提算法实验实现了岩石显微图像的超分辨率处理，并采用客观评价指标对生成的超分辨率岩石图像的性能进行了评价。

1 模型介绍

1.1 生成对抗网络

生成对抗网络由生成网络和判别网络组成，低分辨率图像通过生成网络得到超分辨率图像，判别网络对超分辨率图像和高分辨率图像进行判断。在训练过程中，生成网络和判别网络的对抗过程如下式所示：

(1)

其中：x表示真实数据，Pdata表示其概率分布；z表示生成器的噪声输入，Pz为其相应的概率分布。在生成对抗网络中，要计算的纳什均衡点正是使生成网络和判别网络各自损失函数最小的情况。在上述公式的基础上，可以分别求得理想的生成网络G*和判别网络D*：

(2)

(3)

对于上述的D*，假设生成网络G是固定的，推导如下：

V=Ex～PdatalogD(x)+Ex～Pzlog(1-D(G(z)))

(4)

希望通过寻找一个D使得V最大，则应对于积分中的项f(x)=Pdata(x)logD(x)+Pg(x)log(1-D(x))，无论x取何值都能最大。其中Pdata是固定的，之前假设了G是固定的，所以Pg也是固定的。通过求导可得：

(5)

(6)

由于JS散度是非负的，当且仅当Pdata=Pg时，可以取得全局最小值。故所求的最优生成网络是使得G*的分布满足Pg=Pdata。

1.2 自注意力机制

自注意力机制的使用促使自然语言处理取得了很大的进步，同时上下文信息(context information)对于诸如语义分割、目标检测等视觉任务至关重要。自注意力机制通过三元组提供了一种有效获取更大感受野、捕捉全局上下文信息的建模方式(如图1所示)，推动了视觉任务的发展。

图1 自注意力机制模块图

(7)

2 改进方法

2.1 算法原理

在SRGAN中，通过残差网络提取特征，该网络在提取特征时对于部分边缘信息不敏感，由于岩石显微图像边缘复杂，在提取特征过程中更容易丢失特征信息。本文在SRGAN的基础上，引入了注意力机制模块，希望通过注意力机制模块加强对于边缘信息的描述，并针对岩石显微图像调整了损失函数。

2.2 自注意力机制

本文采用自注意力机制模块加强纹理特征和边缘特征，自注意力机制模块与生成网络的结合如图2所示。自注意力机制模块接收来自上层结构的输出特征图，在模块内部经过相似度计算得到权重并归一化，最终将权重与特征图加权求和，得到特征图中任意两个位置的全局依赖关系。

图2 自注意力机制模块与Generator的结合

2.3 网络结构

本文生成网络和判别网络如图3所示，生成网络在提取特征时采用了残差网络和自注意力机制模块，放大时采用了亚像素卷积模块；判别网络由卷积、注意力机制和激活函数组成。

图3 生成网络和判别网络结构图

生成网络输入一个3×56×56的岩石图像，先后经过一个卷积PReLU层，5个残差块，1个自注意力机制模块，1个卷积BN层，和log2upscale个亚像素卷积块得到大小为3×224×224的超分辨率图像，其中卷积核大小均为3×3，步长和填充均为1。

判别网络输入图像尺寸为3×224×224，输出为经过卷积BN层和Sigmoid激活函数处理之后的概率值。在卷积过程中仍保持3×3大小的卷积核，同时为了更好地对图像进行判断，在其中也加入了自注意力机制模块。

2.4 损失函数

岩石CT图像通过灰度值对岩石内部成分进行了表征，可以对岩体和孔隙进行区分。针对工业CT扫描得到的岩石CT图像，采用图像边缘检测的方法，可以对CT图像进行分割，得到岩石中固体成分和孔隙部分的分类，统计每一类的像素数量即可得到孔隙度。采用灰度直方图来统计图像的灰度分布，当直方图呈现双峰型[11]时，一般选取双峰之间的峰谷所对应的灰度值作为阈值，从而实现岩石固体成分和空隙部分的区分。图4是碳酸盐的灰度直方图，通过分析该图，可以发现对该样本而言，应该选择灰度值85作为岩石固体成分和空隙部分的区分指标，且所得到的孔隙度符合国家给定的碳酸盐孔隙度范围，侧面验证了该处理方法的可行性。

图4 碳酸盐灰度直方图

用psolid和ppore表示岩石固体成分像素数和孔隙像素数，对第n幅CT切片图像I进行分析，得到样品的孔隙度φ，即

φ=ppore(I)/(ppore(I)+psolid(I))。

(8)

(9)

(10)

φi,j(GθG(ILR))x,y)2；

(11)

(12)

lporosity=(φ(ISR)-φ(IHR))2；

(13)

((ISR)x+1,y-(ISR)x,y)2)。

(14)

其中：λ和η为权重参数；r是放大倍数；W是图像宽；H是图像高；IHR表示高分辨率图像；ISR表示超分辨率图像；ILR表示低分辨率图像；θG表示生成网络参数；θD表示判别网络参数；Wi,j,Hi,j分别表示VGG网络中某一层特征图的宽和高；φi,j是VGG网络中某一层的计算结果。

3 实验结构及分析

3.1 实验环境、数据集和评价指标

本文算法在GPU计算型GN10X，8核40GB显存的环境下，通过pytorch框架实现。数据集采用新南威尔士大学的数字岩心超分辨率数据集DRSRD1_2D中的碳酸盐岩石图像。该数据集提供了经过处理之后的灰度数字岩石图像，专门用于图像超分辨率处理。

图像超分辨率处理问题主要研究放大倍数为2和4两种情况，且放大倍数为2的情况在传统神经网络上已经有了很好的效果，因此本文主要研究放大倍数为4的情况。

随机裁剪可以增加数据，还可以弱化数据噪声和增加模型稳定性，于是本文在数据集中对训练集中每张图像进行随机裁剪，获得224×224的图片作为高分辨率图像，通过对数据集中低分辨率图片进行相同操作得到与之对应的56×56分辨率图像。

在超分辨率图像重建的研究中，恢复图像质量的评价指标通常分为主观评价和客观评价。由于主观评价通常容易受到参与人员个人因素的影响，本文采用常用客观评价标准峰值信噪比(Peak Signal to Noise Ratio，PSNR)和结构相似度(Structural Similarity Index，SSIM)进行量化。

3.2 实验结果与分析

3.2.1 岩石孔隙度对结果的影响

本文把岩石孔隙度作为一种约束条件。为衡量其效果，实验通过对放大倍数为4的图像进行了客观指标的评价，加入lporosity的SRGAN与SRGAN在相同实验环境下PSNR和SSIM的对比如图5所示。

图5 SRGAN和加入lporosity的SRGAN对于PSNR和SSIM的对比

通过图5可以发现，增加了损失函数lporosity作为约束条件可以在保证SSIM稳定的前提下，使PSNR的波动更加缓慢，减少了PSNR的突变情况，同时可以发现PSNR可以更快地趋于稳定，加速了网络的收敛过程。

通过对测试集上全体数据不同情况下孔隙度进行分析可得：在加入lporosity作为约束条件之后，孔隙度较SRGAN平均上涨了0.5%左右，更加接近原图的值。在超分辨处理之后能使其孔隙度与原始图片相接近，减少了超分辨处理过程中随机生成造成的影响和处理前后岩石物性特征的变化，对于成矿流体的研究有着极其重要的意义。

随机采样得到5张测试集图片，对其不同情况下的孔隙度进行具体展示，见表1。

表1 原始图像、SRGAN和加入lporosity的SRGAN的孔隙度对比

3.2.2 自注意力机制对结果的影响

为衡量本文所加入的自注意力机制的效果，对放大倍数为4的图像进行了客观指标的评价，并与具有代表性的几种SR方法在同一数据集上进行了对比，结果见表2。

由表2可知，通过在SRGAN中增加注意力机制模块，可以使其获得更高的PSNR值，得到的图像与原图相比失真更小。

表2 测试集DRSRD1_2D中的对比实验结果

通过图6中的SRGAN和本文模型重建后视觉效果的对比，可以发现，在增加了注意力机制之后，对于一些细小的突变位置，SRGAN往往会恢复得相对圆滑，而本文模型对于这些细小的突变则更加敏感，同时在图片的整体感知方面，本文模型和SRGAN效果相近。因此本文在高频特征的恢复上有一定的优势，恢复出来的图像更加接近原始图片。

图6 各算法重建效果比较

3.2.3 3种算法的定量对比

为定量分析加入lporosity和自注意力机制对时间、空间的影响，取得了训练过程中一个epoch所用时间和空间大小，见表3。

表3 一轮训练中SRGAN、porosity-SRGAN和引入自注意力机制SRGAN在时间、空间和模型大小上的对比

通过分析可以发现，加入lporosity会增加其运行时间用于孔隙度的计算，引入自注意力机制则会增加空间用于构建自注意力映射表。自注意力映射表的大小由于受到图像本身大小的影响，在训练和测试时会占用较大的显存空间，但由于其在运行中产生，并不影响模型大小。

4 结 论

本文针对SRGAN超分辨率网络在岩石图像一些细节的恢复上过于圆滑的问题，引入了注意力通道机制，并通过在损失函数中引入岩石图像所特有的孔隙度约束条件，实现了在加速GAN网络训练过程的同时，进一步提升超分辨率图像的视觉效果。

在所采用的数据集上，本文所采用的方法得到的PSNR相比于SRGAN和EDSR分别提升了0.89 dB和0.63 dB。然而由于超分辨率重构本质上是一个病态问题，找到合适的高低分辨率图像映射关系极其困难，因此与原图相比仍有提升的空间。由于注意力机制的引入导致模型训练和测试时会占用大量的显存，同时也会消耗较多的时间，故本文下一步将着手研究在保障重构质量的前提下提升算法的训练、测试效率。