2020-05-20 图像去噪的超核神经网络结构搜索

Abstract

神经架构搜索(NAS)的最新进展导致在没有实质性人工监督的情况下,为图像分类、目标检测或语义分割等任务找到了最新的人工神经网络(ANN)解决方案。本文主要针对图像去噪这一稠密预测任务,对神经网络进行研究。由于昂贵的训练过程,大多数用于图像增强的NAS解决方案依赖于强化学习或进化算法探索,通常需要数周(甚至数月)的训练。因此,我们引入了一种新的高效的超级核技术实现,它能够快速(6-8RTX2080GPU小时)对模型进行单次密集预测训练。我们证明了我们的方法在SIDD+图像去噪基准上的有效性

Introduction

神经网络结构搜索(NAS)寻求自动选择最佳神经网络结构的过程。由于模型结构选择的可训练公式,它不仅可以优化隐藏任务,还可以优化其他执行属性,例如内存约束或推理时间。已经提出了许多搜索技术,包括基于增强学习的搜索技术[41]、进化算法[29]或梯度下降的搜索技术[22、36、20、5]。

大多数NAS算法所面临的最大挑战是在搜索时间和搜索过程的内存需求之间进行权衡。在经典的方法中,即强化学习或进化算法中,来自搜索空间的多个样本在并行训练过程中收集的统计数据驱动着探索过程。由于探索算法通常需要从头开始训练这些样本,搜索过程需要数百[21]或数千[42,28,6]个GPU小时才能形成。为了缓解这一问题,人们提出了多种技术。其中最流行的一种方法要么缩小搜索范围以找到一个本地单元结构,该结构随后在整个网络中复制[42],要么在经过训练的模型样本中共享权重[27]。必须注意,给定训练过程的内存需求与探索的搜索空间中消耗内存最多的模型成正比。这种特性使得这些方法成为执行图像增强任务的神经架构搜索的良好开端,在这些任务中,单个模型的大小常常填满整个GPU内存。

one-shot :另一种方法,通常称为单点搜索[3,4,15],寻求上述问题的不同解决方案。它以超级网络的形式对搜索空间进行建模,超级网络本身也是一个神经网络。然后使用梯度去嗅进行搜索过程。这些方法经常将整个搜索时间从几个星期缩短到几个小时[32]。不幸的是,这种方法的一个巨大的缺点是,由于超级网络包含了整个搜索空间,它的内存需求远远大于搜索空间中典型网络的内存需求。对于密集预测任务来说,这个问题更为重要,因为即使是单个模型在内存使用方面也可能要求很高。

一个有效的中间点是超核方法。它将结构上的超级网络部分简化为一个卷积核,从而将其内存需求减少了一个数量级。这种改进是以将搜索空间选择缩小到核大小和标准卷积算子的滤波器数目为代价的。

超长训练时间和内存使用的一个常见解决方案是使用代理数据集。该数据集通常较小,无论是在示例数量还是图像大小方面(例如,使用CIFAR10而不是ImageNet进行图像分类)。在对较小的数据集进行初始训练之后,搜索过程将最终的体系结构传输到目标任务。有几种方法试图直接在目标数据上以proxyless的方式执行架构搜索过程。这可以通过使用简化的搜索空间[36]或在训练期间压缩/修剪超级网[5]来实现。无代理的方法对于图像增强任务(如图像去噪)可能特别有吸引力。对于这些任务,使用代理数据集是困难的,因为正确的训练过程需要高分辨率的图像。

在这项工作中,我们提出了一种新的松弛超核图像去噪解决方案,它具有训练速度快(6-8gpu小时),内存效率高(一个超级网络适合一个批量大小为4的GPU),并且可以使用128x128的输入图像分辨率以无障碍的方式进行训练。我们在SIDD+[11]上评估我们的模型(2020年真实图像去噪挑战的一部分[2,11,1])数据集,以实现图像去噪的最新结果。

2.Related work


2.1 Reinforcement Learning and Evolutionary algorithms for Neural architecture search

最初的方法使用强化学习和进化算法来处理神经架构搜索。一篇开创性的论文[41]使用了强化学习控制器来产生网络结构,然后从头开始训练和评估。正因为如此,整个培训过程花费了数千GPU小时来为CIFAR10数据集培训最先进的模型。为了加快计算速度,在接下来的工作[42]中,搜索空间被缩小到只有少量的子网络结构,随后重复以形成最终的体系结构。这种方法后来通过在不同的子网之间共享权重来改进,以便传递知识[27]。类似地,[28]使用遗传算法而不是强化学习来驱动优化过程,而[6]结合了这两种方法。

由于这些方法的内存需求与搜索空间中最大模型的记忆需求成正比,因此该技术被用于超分辨率图像增强任务,如医学图像去噪[23]、图像恢复[37、34、16]和图像修复[19]。

2.2 Single-shot approaches

另一种方法是以神经网络的形式对搜索空间本身进行建模。通常,它是通过一个连续的采样松弛过程来完成的,该过程以可微的方式逼近离散的建筑选择过程。e、 g.在[22]中,作者使用softmax权重来近似操作选择。更准确地说,ap代理是通过使用softmax权重和可学习登录来组合每个操作的输出来实现的。采用基于两级梯度的优化过程对超网络进行训练,将卷积权值和结构参数的训练分离开来。最终的架构是通过选择最高logit值的操作获得的。

[20]的作者进一步推动了这种方法,并且使用softmax权重作为全局CNN体系结构(频道、步幅、深度、连接性)和本地单元结构的松弛。该算法为语义图像分割任务找到一个网络,该任务根据提供的数据集调整其结构。搜索空间形成由称为单元的本地层结构构建的网格。每个单元是一个一般的无环图,k个单独选择和优化的操作之间有可学习的连接。网格由一个复制单元和一个连接结构组成,它支持多种流行的CNN设计,如DeepLabv3、AutoEncoder或堆叠沙漏。架构搜索阶段在P100 GPU上持续几天。

考虑到当模型分布的熵很高时,softmax权重可能是离散采样的不精确近似,[36]将Gumbel softmax近似引入神经架构搜索[17]。在这种技术中,用于组合不同操作的随机权重似乎比softmax更像一个热系数。提出的解决方案是一种计算量小的算法,它为图像分类任务找到一个设备感知的CNN结构。该算法在预定义的全局CNN结构上工作,即频道、步幅和深度。搜索空间由一组张量操作跨越,包括不同的卷积设置、池和跳过连接。通过对每一层分别进行梯度下降,在考虑分类指标和触发器的情况下,找到最优操作。该算法为不同的硬件设置(如三星S10和iPhone X)找到不同的体系结构。同样,体系结构搜索阶段在数百GPU小时内非常耗时。

上述单发方法的巨大缺点是,它们需要计算所有可能的操作,以便执行搜索的单个迭代。在文献[5]中,作者试图通过引入搜索空间子抽样来规避这一问题,从而降低存储成本。在每次迭代中,仅对随机选择的可能操作子集执行训练更新。尽管如此,这些模型通常需要的内存大约是搜索空间中要求最高的模型的两倍。

2.3 Superkernels

为了缓解内存问题,[32]引入了一种称为超级内核的新技术。该方法将搜索过程从操作选择切换到核选择。它使用一个具体的分布,以便将最大核的一部分归零,该最大核在搜索空间的所有模型中进行训练和共享。由于卷积运算输出的大小通常比其参数的大小大几个数量级,因此在核空间中执行搜索的额外成本可以忽略不计。在第四章中,我们介绍了一些基于Gumbel-Softmax松弛的新技术。

2.4 Deep Learning for Image denoising

自从第一个成功的MLP方法[30]取得了与当时最先进的BM3D方法[9]相当的性能以来,几种深度学习方法被应用于图像去噪任务。在文献[38]中,作者使用了叠置的卷积块来逼近噪声图像与其清理后的图像之间的残差。类似地,[25]不仅在网络的输入和输出之间,而且在模型链卷积块之间引入了剩余连接。在Memnet[33]中,作者引入了基于递归单元和选通单元的记忆机制,以结合输入图像的多级表示。

上述所有方法都不使用任何形式的池/下采样,这会带来很高的计算成本。为了缓解这个问题[24]使用了U-NET架构。通过应用池/上采样操作,它们降低了内部循环卷的分辨率。它大大减少了计算的内存和计算负担。在SGN[13]中,作者注意到,由于大量低层图像信息丢失,使用传统的池/上采样操作显著降低了重建质量。正因为如此,他们引入了一种洗牌(亚像素)重缩放,旨在保留更多细节,而不会显著增加计算成本。此外,由于自上而下的自引导算法,该算法能够实现多尺度图像特征的轻量化组合。

3.Models for image denoising

在本节中,我们将介绍我们在实验中使用的基本体系结构。我们通过引入超核而不是标准卷积滤波器,将它们转化为神经结构搜索超网络。

3.1. Superkernel-based Multi Attentional Residual U-Net

基于superKernal的多注意力的残差U-Net网络的架构如图1所示。该网络由多个子网络SK-A-RES-UNET同时训练而成。每个子网的输出被传递给信道注意块[39]。作为最后一层,我们使用核大小为k=3的卷积来减少信道数。最终输出与输入图像相加。

基于超核的注意剩余U-Net子网络结构如图2所示。该网络是对U-Net体系结构的改进,增加了注意机制和基于超核的密集连接剩余块(SK-DCRB)。每个编码器层具有相同的结构。具体来说,它由ReLU激活函数和SK-DCR块构成卷积层,核大小k=3。与普通U-Net架构一样,编码器层的输出在同一空间层传递给解码器层。每个解码器层接收跳过连接的输出和前一层的结果。与U-Net体系结构不同,在每个解码器层中进行级联操作后,应用信道注意块(CAB)。我们使用步长为2的卷积进行下采样,使用洗牌/亚像素层[13]进行上采样。

图3给出了超核密集连接剩余块的模式。该区块具有[26]中提出的类似结构。每个SK-DCR块包含三个卷积层,然后是ReLU激活函数。我们用超核替换了前两个卷积运算。这种方法允许网络学习最合适的内核大小和增长率。输出卷积层恢复滤波器的数目。

3.2. Superkernel SkipInit Residual U-net

网络的架构与图2中描述的架构类似。不同的部分是实际的卷积块。在基于超核的SkipInit剩余U-net中,有多个(在所选网络中有2个)密集连接的剩余块(例如参见[40])与SkipInit标量乘法一起稳定学习,如[10]所述。图4显示了该块的轮廓。简单地说,每个剩余分支的输出都是可训练标量的倍数,该标量初始化为零,从而使块在开始时像身份转换一样工作。如果优化选择使用剩余分支,那么这个标量(在图4中由alfa表示)将改变。特定级别的筛选器数是其正上方级别的筛选器数的两倍,从级别0的64开始。卷积超参数的训练遵循第4.1节描述的联合超核机制。

4. Superkernels

在每一种结构中,我们应用神经结构搜索,以找到每个卷积块的最佳滤波器数目和核大小。搜索过程基于superk-ernel技术的多次修改,即:

Factorized Gumbel Superkernel,

Joint Gumbel Superkernel,

Filterwise Gumbel Superkernel,

Filterwise Attention-based Superkernel.

在每一种方法中,神经结构搜索给出了内核大小和过滤器数量的不同可能性。每个选择都是通过对一个称为超级核的最大核进行适当的切片来获得的,并且这种切片的执行方式在不同的方法中是不同的。在每种情况下,算法的目标是优化一组可能切片上的结构分布。我们稍后将这些结构参数的优化过程称为搜索。下面可以找到每个NAS方法的详细描述。

4.1. Joint Superkernel

使用这种技术,在一组切片Sk上进行搜索,其中k∈{k1,…knk}是一个核的大小,f∈{f1,…,fnf}是多个滤波器。fi和ki的值在不同的模型和架构中都有所不同。对于每一个超核,训练并共享核大小为maxi=1,…nk ki和maxi=1,…nf-fi滤波器数的最大超核。slice Sk,f由一个中心子核(核大小为k)和最大超核的第一个f过滤器组成(见图5)。

该结构分布被建模为一组所有可能元组(k,f):k∈{k1,…,knk},f∈{f1,…,fnf}上的softmax极大分布,每个元组有一个logit参数θ(k,f)

4.2. Factorized Superkernel

在这项技术中,子核集与第4.1节中的联合超核相同,但结构分布被分解为两个独立分布pk,f=pk-pf,其中pk和pf是一组可能的核和滤波器大小上的softmax分布。在这种情况下,优化后的参数形成两组logit参数:{θki:ki∈{1。核切片和{θfi:fi∈{1。过滤片的nf}。与联合超核相比,这种因子分解以较小的建模灵活性为代价,显著减少了结构分布(从nf-nk到nf+nk)的参数数目。

4.3. Filterwise Superkernel

在这种技术中,搜索是在一组切片Sk上进行的,其中k∈{k1,…,kn}-是一个核的大小,M∈{0,1}F是F势滤波器的所有可能子集的集合。对每个子核训练一个核大小为max-ki和F的最大超核,并在每个子核上共享。slice Sk,M由一个内核大小为k的中心子内核和那些Mi=1的过滤器组成(见图6)。

将结构分布分解为两个独立的分布pk,M=pk pM,其中pk是可能核集合上的softmax分布,pM由F-i.d.Bernoulli分布建模。pk分布是由一组logit参数{θki:i∈{1。,kn},其中Mi上的每个F独立Bernoulli滤波器分布由单个logit参数θfi控制。

4.4. Filterwise Attention-based Superkernel

在这种技术中,子核的集合与在Filterwise超核中相同(秒)。4.3节)。同样,核大小k上的分布也是可能核集上的softmax分布。一个实质性的区别在于掩模M∈{0,1}F上的分布是受管辖。一次同样,maskM上的每个分布都是独立的Bernoulli。然而,这一次,这些分布中的每一个都由单个基logit参数θb和经由∈Rl的fi注意密钥向量参数化。贝努利分布θfi的最终logit是根据注意机制计算的[35]:

4.5. Details of superkernel implementation

4.5.1 Sampling Relaxation 抽样松弛

由于切片是一种非连续的操作,因此不能用梯度下降法直接优化网络的结构参数。为了实现这类训练,我们应用了采样的连续近似,即Gumbel-Softmax和松弛Bernoulli分布[17]。对于因子化超核,我们给出了这种松弛的一个概要。

我们可以用采样的子内核Sk,f在卷x上重写卷积,方法如下:

其中Sk′,f′通过所有可能的超核集合,I是一个指示函数。Gumbel松弛相当于以下近似值:

其中,GS是来自Gumbel-Softmax分布的样本,其中θ(k,f)logit参数化了一组可能的切片元组上的Softmax分布(见第4.1节)。

4.5.2 Mask sampling reparametrization

在上述近似中,仍需要计算nknf操作才能执行最终求和。我们决定使用掩模采样技巧来减少计算和内存负担。超核切片可以是重参数化的Sk,f=S*Ik,f其中S是每个子核共享的超核,Ik,f是应用于与S形状相同的所有核的适当子核切片,*是两个张量的Hadamard积。此重新参数化允许对采样操作进行以下重新设置:

上述方程的RHS只需一次卷积运算,最大超核S用平均子核来做mask。这大大减少了计算负担,无论是在FLOPs和内存需求方面,并使图像去噪的全NAS处理成为可能。

4.5.3 Mask Sampling reparametrization - issue with a non-linear activation function 

4.5.3掩模采样重新参数化-非线性激活函数问题

人们可能注意到,对于不仅涉及Conv算子而且还涉及非线性激活函数的搜索过程,上述推导是不准确的。这是因为这种关系:

仅适用于加性函数f。在我们的实验中使用的ReLU和PReLU都不具有此特性。因此,我们测试了两种可能性:

full - 为了保持计算的可行性,我们决定忽略这个事实,把上面方程的RHS当作LHS的近似值,

separate - LHS所有的左半部分都单独计算的。

我们在一组较小的架构上测试了这两种方法,没有发现任何显著的差异。如果一个模型的名称中有一个单独的组件,则表示它是使用单独的方法训练的。否则就要接受full。

4.5.4 Bias sampling 偏差采样

我们已经应用了一种类似于卷积核的采样技术来处理偏差。当偏置向量的大小等于滤波器的数目时,通过平均除滤波器维数外的所有维数的核掩模,可以得到合适的偏置切片掩模。

4.5.5 Final model distillation 最终模型蒸馏

一旦训练了模型的结构参数和权值,就需要提取出具有最佳核和滤波器尺寸的模型。我们采用了以下蒸馏策略:

对于由softmax分布控制的结构参数,选择了概率最大的选项,

对于由Bernoulli分布控制的结构掩码,如果其logit大于0.5,则选择过滤器。

这个蒸馏过程是为每个超核独立完成的。我们计划在未来的工作中使用更复杂的蒸馏策略。

6. Discussion

在所有介绍的架构中,至少有一种神经架构搜索方法比无NAS基线有显著改进。最佳单一模型-52.74psnr使用分离联合超核-并通过SkipInit剩余U-网实现模型1。对于小型模型(单个子网),基线架构和搜索架构之间的差异最为显著(约为PSNR的0.12),这从移动设备上模型部署的角度来看是有希望的。

消融研究表明了自我强化技术的重要性。平均而言,它在PSNR(0.21)和SSIM(0.0002)指标上都有改进。有趣的是,在没有自我意识的情况下,NAS模型的性能比其无NAS基线的性能要差得多。

7. Potential improvements and future work

尽管我们的搜索过程对小模型的PSNR改善最为显著,但这是以选择具有更大核尺寸和更高滤波器数量的卷积核为代价的。然而,内核大小的可微表示允许在优化过程中引入额外的计算约束。我们期望,配备额外计算正则化的训练将导致找到可部署在便携式设备上的模型。

我们的搜索过程类似于贝叶斯深度学习的变分推理方法;因此,我们的方法可能会受到[12]中所报告的搜索空间探索不足的影响。这个问题可能是因为我们的方法只探索接近梯度下降轨迹的点,这会缩小搜索空间的探索区域,使其对权值初始化具有潜在的敏感性。我们计划在未来的工作中探索这种敏感性。此外,在我们的训练过程中,权重和结构参数都被联合优化。我们计划探索两级优化(类似于[22]),以检查这些参数的耦合如何影响结果。

由于不同超核内的采样是独立执行的,因此可能会出现与公平性相关的问题[8]。因此,根据马太定律,网络可能很快收敛到随机次优解。我们计划测试更高级的概率模型(例如基于马尔可夫随机场、递归神经网络或注意力),以便对更高级的搜索空间建模并进行更广泛的探索。

此外,我们观察到训练过程可能会对最终结果产生重大影响。特别是,与非超核模型相比,配备超核的模型可能需要更长的训练时间(更高的耐心)。我们怀疑这可能是NAS机型在没有自我意识的情况下性能下降的原因。我们计划在今后的工作中研究这个问题。

8. Conclusion

在这项工作中,我们介绍了一种快速且轻量级的用于神经架构搜索的图像去噪算法。我们证明,它可以达到最先进的结果,不加NAS解决方案的显著幅度。我们已经证明,所提出的超核技术可以在几小时内达到与最先进的图像去噪架构相当的效果。

你可能感兴趣的:(2020-05-20 图像去噪的超核神经网络结构搜索)