《Superkernel Neural Architecture Search for Image Denoising》

神经体系结构搜索（NAS）的最新进展导致找到了新的最新人工神经网络（ANN）解决方案，无需大量人工监督即可完成图像分类，对象检测或语义分割等任务。 在本文中，我们专注于探索NAS以进行图像降噪的密集预测任务。 由于昂贵的培训程序，大多数用于图像增强的NAS解决方案都依赖于强化学习或进化算法探索，通常需要数周（甚至数月）的时间进行培训。 因此，我们引入了各种超级内核技术的新有效实现方式，这些技术可对模型进行快速（6-8个RTX2080 GPU小时）单次训练以进行密集的预测。 我们在SIDD +基准上证明了我们的方法对图像去噪的有效性[11]。

1 介绍

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

大多数NAS算法所带来的最大挑战是搜索时间与搜索过程的内存需求之间的权衡。 在经典方法（即强化学习或进化算法）中，从架构空间并行训练多个样本期间收集的统计信息驱动了探索过程。由于探索算法通常需要从头训练这些样本，因此搜索过程需要执行数百[21]或数千[42、28、6]个GPU小时。 为了减轻这个问题，提出了多种技术。 最受欢迎的一种方法是缩小搜索范围以找到局部细胞结构，该结构随后将在整个网络中复制[42]，或者在经过训练的模型样本之间共享权重[27]。 至关重要的是要注意，给定训练过程的内存需求与所探索的搜索空间中最消耗内存的模型成比例。 此属性使这些方法成为执行神经架构搜索以执行图像增强任务的良好第一步，其中单个模型的大小通常会填满整个GPU内存。

一种通常称为单发[3，4，15]搜索的替代方法，针对上述问题寻求了不同的解决方案。 它以超级网络的形式对搜索空间建模，超级网络本身也是神经网络。然后使用梯度下降执行搜索过程。 这些方法通常将整体搜索时间从几周减少到几个小时[32]。 不幸的是，这种方法的一个巨大缺点是，由于超级网络包含整个搜索空间，因此它的内存需求比搜索空间中的典型网络要大得多。 对于密集的预测任务而言，此问题甚至更为重要，因为即使在使用内存的情况下，即使是单个模型也可能要求很高。

一个有效的中点是超级内核方法。它将结构性超级网络部分简化为单个卷积内核，从而将其内存需求降低了一个数量级。 以将搜索空间选择范围缩小到内核大小和标准卷积运算符的过滤器数量为代价来实现此改进。

延长训练时间和内存使用量的常见解决方案是使用代理数据集。 就示例数量和图像大小而言，此数据集通常较小（例如，使用CIFAR10代替ImageNet进行图像分类）。 在对较小的数据集进行初始训练之后，搜索过程会将最终的体系结构转移到目标任务。 几种方法尝试以无代理方式直接在目标数据上执行体系结构搜索过程。 这可以通过使用simpliarXiv：2004.08870v1 [eess.IV] 2020年4月19日限定搜索空间[36]或通过在训练过程中压缩/修剪超网来实现[5]。 无代理方法对于图像增强任务（例如图像去噪）可能特别有吸引力。 对于这些任务，很难使用代理数据集，因为正确的训练过程需要高分辨率的图像。

在这项工作中，我们提出了一种用于图像去噪的新型轻松超级内核解决方案，该解决方案训练速度快（6-8个GPU小时），内存效率高（一个超级网络适合批处理大小为4的单个GPU），并且可以在 使用128x128的输入图像分辨率的无代理方式。 我们在SIDD + [11]（NTIRE 2020真实图像去噪挑战[2，11，1]的一部分）数据集上评估模型，以实现最新的图像去噪。

2 相关工作

2.1 神经架构搜索的强化学习和进化算法

最初的方法通过使用强化学习和进化算法解决了神经体系结构搜索。
   一篇开创性的论文[41]使用增强学习控制器来生成网络结构，然后从头开始对其进行训练并进行评估。 因此，整个训练过程花费了数千个GPU小时，以训练出CIFAR10数据集的最新模型。 为了加快计算速度，在以下工作中[42]，将搜索空间缩小到仅少数子网结构，随后重复进行以形成最终的体系结构。 该方法后来通过在不同子网之间共享权重进行了改进，以传输知识[27]。 同样，[28]使用遗传算法代替强化学习来驱动优化过程，而[6]则将这两种方法结合在一起。

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

2.2 单发方法

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

[20]的作者进一步推动了该方法的发展，并使用softmax权重作为全局CNN架构（通道，步幅，深度，连通性）和局部小区结构的放松。 该算法找到用于语义图像分割任务的网络，该网络将其体系结构调整为提供的数据集。 搜索空间形成由称为单元的局部层结构构成的网格。 每个单元都是一个普通的非循环图，在k个单独选择和优化的操作之间具有可学习的联系。 网格由复制的单元格和连接结构组成，可以对多个流行的CNN设计进行建模，例如DeepLabv3，AutoEncoder或Stacked Hourglass。 在P100 GPU上，架构搜索阶段持续数天。

假设当模型分布的熵高时，softmax权重可能是离散采样的不精确近似值，[36]将Gumbel Softmax近似值引入了神经体系结构搜索[17]。 在这种技术中，用于组合不同操作的随机权重似乎比softmax更像单热系数。 提出的解决方案是一种计算轻量级算法，可找到用于图像分类任务的设备感知CNN架构。 该算法适用于预定义的全局CNN结构，/即通道，步幅和深度。 搜索空间由一组张量操作扩展，包括不同的卷积设置，池化和跳过连接。 考虑到分类指标和FLOP，可以通过梯度下降分别找到每个层的最佳操作。 该算法为不同的硬件设置（即三星S10和iPhone X）找到了不同的架构。同样，架构搜索阶段要花费数百个GPU小时。

上面介绍的单发方法的巨大缺点是，它们需要计算所有可能的操作才能执行搜索的单次迭代。 在[5]中，作者试图通过引入搜索空间二次采样来规避这一问题，从而降低了存储成本。 在每次迭代中，仅对随机选择的可能操作子集执行训练更新。 尽管如此，这些模型通常仍需要比搜索空间中最苛刻的模型大约两倍的内存。

2.3 超级内核

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

2.4 深度学习进行图像降噪

自第一个成功的MLP方法[30]获得的性能可与当时最先进的BM3D [9]方法相比，将几种深度学习方法应用于图像去噪任务。 在[38]中，作者使用了堆叠的卷积块，以近似嘈杂图像与其清洁版本之间的残差。 类似地，[25]不仅在网络的输入和输出之间引入了残余连接，还在模型链式卷积模块之间引入了残余连接。 在Memnet [33]中，作者介绍了基于递归和门控单元的存储机制，以便组合输入图像的多级表示。

上面介绍的所有方法都没有使用任何形式的合并/下采样形式，这会导致较高的计算成本。 为了减轻这个问题，[24]使用了U-NET架构。 通过应用合并/上采样操作，它们降低了内部卷积量的分辨率。 它显着减少了计算的内存和计算负担。 在SGN [13]中，作者注意到使用传统的合并/上采样操作会大大降低重建质量，因为会丢失许多低级图像信息。 因此，他们引入了改组（子像素）缩放功能，旨在保留更多细节，而不会显着增加计算成本。 此外，由于采用了自上而下的自导算法，因此可以轻量化多尺度图像特征的组合。

3. 图像去噪模型

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

3.1 基于超级内核的多注意残量U-Net

图1中显示了基于超级内核的多注意残差U-Net网络的建议体系结构。该网络包含多个同时被训练的名为SK-A-RES-UNET的子网。 每个子网的结果传递到频道关注模块[39]。 作为最后一层，我们使用内核大小为k = 3的卷积来减少通道数。 最终输出将另外与输入图像相加。

基于超级内核的注意力残留U-Net子网的体系结构如图2所示。该网络是对U-Net体系结构的修改，另外还配备了注意机制和基于超级内核的密集连接的残余块（SK-DCRB）。 每个编码器层具有相同的结构。 具体来说，它包括一个内核大小为k = 3的卷积层，其后是ReLU激活函数和SK-DCR块。与普通U-Net体系结构一样，编码器层的输出在相同的空间级别传递到解码器层。每个解码器层都从跳过连接获取输出，并从上一层获取结果。 与常见的U-Net体系结构不同，在每个解码器层中进行级联操作之后，将应用通道注意块（CAB）。 我们将步长为2的卷积用于下采样，将混洗/子像素层[13]用于上采样。

图3给出了Superkernel密集连接残差块的架构。 该块的结构与[26]中提出的类似。 每个SK-DCR块均包含三个卷积层，后跟ReLU激活功能。 我们用超级内核代替了前两个卷积运算。 这种方法允许网络学习最合适的内核大小和增长率。 输出卷积层恢复滤波器的数量。

3.2 超级内核SkipInit残留U-net

网络的体系结构与图2所示的体系结构相似。不同之处在于实际的卷积块。 在基于Superkernel的SkipInit残余U-net中，有多（在选定的网络中为2个）密集连接的残余块（例如，参见[40]），其具有SkipInit标量乘法以稳定学习，如[10]中所述。 该块的轮廓如图4所示。简而言之，每个残差分支的输出乘以可训练的标量，该可训练的标量被初始化为零，从而使该块在开始时就像身份转换一样工作。 如果优化选择使用残差分支，则此标量（在图4中用alfa表示）将改变。 特定级别的过滤器数量是其正上方的过滤器数量的两倍，从0级别的64开始。卷积超参数的训练遵循4.1节中描述的“联合超级内核”机制。

4. 超核

在这些架构的每一个中，我们都应用了神经架构搜索，以便为每个卷积块找到最佳的滤波器数量和内核大小。 搜索过程基于超级核技术的多种修改，即：

•因式分解的Gumbel超级核，

•联合Gumbel超级核，

•Filterwise的Gumbel超级核，

•Filterwise的基于注意的超级核。

在每种方法中，神经体系结构搜索都探索了内核大小和过滤器数量的不同可能性。 每个选择都是通过对最大内核（称为超级内核）进行适当切片而获得的，并且在不同方法之间，执行切片的方式也有所不同。 在每种算法中，算法旨在优化一组可能的切片上的结构分布。 我们稍后将参考。

4.1 Joint Superkernel

使用该技术，在一组切片上执行搜索，其中是一个内核的大小，是许多过滤器。 值fi和ki在不同的模型和体系结构上都不同。 对于这些超级内核中的每一个，将训练内核大小为和的最大超级内核，并在每个子内核之间共享它们。 切片Sk; f包含一个内核大小为k的居中子内核和最大超级内核的前f个过滤器（请参见图5）。

结构分布被建模为所有可能的元组上的softmax分布： 每个元组具有单个logit参数。

4.2 分解超级内核

在该技术中，子内核集与第4.1节中的“联合超级内核”相同，但是结构分布被分解为两个独立的分布 其中pk和pf是一组可能的内核和过滤器大小上的softmax分布。 在这种情况下，优化参数形成了两组logit参数：用于内核切片，用于过滤器切片。 与“联合超级核”相比，这种分解显着减少了结构分布的参数数量，但具有较小的建模灵活性。

4.3 Filterwise超级内核

在该技术中，搜索是在一组切片上执行的，其中-是内核大小，势滤波器的所有可能子集的集合。 对于每个子内核，将训练并在每个子内核之间共享内核大小为和F的最大超级内核。 切片包含一个内核大小为k的居中子内核，以及的那些过滤器（请参见图6）.

结构分布分解为两个独立的分布，其中pk是整个可能内核集上的softmax分布，而pM由：Bernoulli分布建模。  pk分布由一组logit参数fθki参数化：，其中Mi上F个独立的Bernoulli过滤器分布均由单个logit参数控制。

4.5 基于筛选器的基于注意的超级内核

在此技术中，子内核集与Filterwise Superkernel（第4.3节）中的相同。 同样，内核大小k上的分布也是可能的内核集上的softmax分布。 掩模的分布方式有很大的不同被管理。再一次，掩模Mi上的每个分布都是独立的伯努利。 然而，这次，这些分布中的每一个都由单个基本logit参数和注意关键向量来参数化。 伯努利分布θfi的最终logit是根据注意力机制计算的[35]：

4.5 超级内核实现的详细信息

4.5.1 采样松弛

由于切片是非连续操作，因此无法使用梯度下降直接优化网络的结构参数。 为了实现这种训练，我们采用了采样的连续近似，即Gumbel Softmax和Relaxed Bernoulli分布[17]。

对于因式分解超级内核的情况，我们给出了这种放松的概述。我们可以按以下方式用样本子核Sk; f重写卷x上的卷积：

其中遍历所有可能的超级内核，而是一个指标函数。  Gumbel松弛等效于以下近似值：

其中GS是来自在一组可能的切片元组上的softmax分布的对数参数化的Gumbel Softmax分布的样本（请参见第4.1节）。

4.5.2  模板采样重新参数化

在上述近似中，为了执行最终求和，仍然需要计算nknf运算。 我们决定通过使用掩码采样技巧来减少计算和内存负担。 超级内核切片可以重新设置为，其中S是每个子内核共享的超级内核，是适用于形状与S相同的所有人内核的适当子内核切片，∗是 Hadamard两个张量的乘积。 通过重新参数化，可以重新配置以下采样操作：

上式的RHS仅需要一次卷积运算，而最大的超级内核S会被平均的子内核掩码所掩盖。 这在FLOP和内存需求方面都大大降低了计算负担，并使NAS能够完全处理图像降噪。

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

可能会注意到，以上推导对于不仅涉及Conv运算符而且涉及非线性激活函数的搜索过程而言并不准确。 这是由于以下关系：

仅对加法函数f成立。 在我们的实验中使用的ReLU和PReLU都不具有此属性。 因此，我们测试了两种可能性：

•完整-为了保持计算的可行性，我们决定忽略这一事实，并将上述方程式的RHS视为LHS的近似值；

•分开-其中LHS的所有组成部分均单独计算。

我们已经在一组较小的体系结构上测试了这两种方法，但没有发现任何显着差异。 如果模型名称中包含单独的组件，则表明模型是使用单独的方法进行训练的。 否则，将对其进行全面对待。

4.5.4 偏差采样

我们将与卷积核类似的采样技术应用于偏差。 由于偏置矢量的大小等于滤波器的数量，因此可以通过对除滤波器尺寸以外的所有尺寸的内核掩码求平均来获得适用于偏置片的掩码。

 4.5.5最终模型蒸馏

训练完模型的结构参数和权重后，我们需要以最佳内核和过滤器尺寸蒸馏模型。 我们应用了以下蒸馏策略：

•对于由softmax分布控制的结构参数，选择了具有最大概率的选项；

•对于由Bernoulli分布控制的结构蒙版，如果对数大于0：5，则选择过滤器。

该蒸馏程序是针对每个超级内核独立完成的。 我们计划在未来的工作中使用更复杂的蒸馏策略。

5. 实验结果

我们在原始RGB图像降噪任务中进行了几次实验，以测试我们搜索最佳内核尺寸/输出通道数组合的方法是否为针对同一任务开发的基准（no-NAS）解决方案提供了明显的好处。 在本节中，我们简要描述了训练过程和使用的模型。 我们还根据获得的PSNR和SSIM分数提供最终结果。

5.1 训练细节

我们在NTIRE2020实像去噪rawRGB竞赛组织者提供的SSID + [11]训练数据上训练了所有模型。 在我们的训练过程中，我们将挑战图像分为两部分：一个训练集的图像占90％，其余10％用于验证。 根据数据集移动电话类型对拆分进行了分层。我们使用比赛组织者提供的验证数据作为最终测试集。 我们将训练和验证图像切成128x128的补丁。 我们使用学习率2 * 10−5的Adam优化器[18]进行模型训练。 在学习过程中，学习率保持恒定。 由于我们的管道包含大量不同结构和能力的模型，因此我们决定使用一种流行的早期停止技术来限制训练时间。 在学习过程中，我们监控了验证数据集上的PSNR值。 我们每隔1K步评估一次模型。 如果在30次验证评估中PSNR没有改善，我们将停止训练过程。 根据所用硬件（GeForce GTX 1080Ti或GeForce RTX 2080Ti）的不同，我们对模型的培训持续了几个小时到三天。

5.2  超级内核SkipInit残差U-net

本节中模型的总体架构基于[26]中介绍的U-net。 本小节中的模型1是第3.2节中描述的模型，其U净深度为3，在DCR块中具有三个可搜索的增长率0.2、0.4和0.6。 每个U-net级别有2个深度为2的DCR块（编码器和解码器部分都使用）。 在级别0处，最大可能的输出通道数为32，并且随着每一层的算术级数增加而增加。 模型2与模型1的区别仅在于U-net深度（2而不是3）和0级的最大通道数（64而不是32）。 这些模型的性能可以在表1和表2中看到。

5.2.1选择的体系结构

两种模型的搜索空间包括DCR块的内核大小（3或5）和增长率（0.2、0.4、0.6）的选择，请参见图4和3.2节。 模型1中有28个DCR模块，模型2中有20个DCR模块。

完整版本的超级内核联合框架中的模型1连续选择了5号内核，只有少数例外。 关于增长率，除三个DCR块外，所有块均选择了最大的增长率。 在单独的版本中，情况类似，内核大小3仅选择一次。

在具有完整版本的超级内核联合的模型2中，该模型具有较高的过滤器基线数量，但就U-net架构而言较浅，因此几乎在所有地方都选择了5号内核。 增长率略有不同，大约40％的DCR块选择增长率为0.4，其余选择0.6。

5.3 基于超级内核的多注意残差U-Net

表3和表4包含第3.1节中描述的体系结构的PSNR和SSIM分数。 我们用不同数量的子网和超级内核类型训练了几种模型。 表3列出了不具有自集合的结果，而表4中显示的模型则采用了自集合技术。 我们使用具有1、3和6个子网的体系结构。 每个子网具有相同的结构，如图2所示。如表4所示，自集成技术的应用显着改善了平均PSNR和SSIM度量。 对于包含六个子网的模型，可以获得最佳的PSNR结果。 配备有六个子网（具有自集成）的模型的平均PSNR值分别比1-子网和3子网模型高0.13和0.03。

5.3.1 选择的架构

该模型的搜索空间包括内核大小和密集连接的残差块的增长率。 对于DCR块中的每两个前两个卷积层，超级内核必须选择k 2 f3中的滤波器大小，增长率。 最终的超级内核选择因网络部分而异。 在编码器部分，超级内核选择了5和7的内核大小以及1和0.5的增长率。 在解码器部分，NAS仅选择了7个内核大小和最大过滤器数（r = 1），只有一些小例外。

 6 讨论

在介绍的所有体系结构中，至少有一种神经体系结构搜索方法相对于noNAS基线有了显着改进。 最佳单一模型-PSNR为52.74的单独联合超级内核-是通过SkipInit残量U-net实现的模型1。基线架构和搜索的架构之间的差异最为明显（PSNR约为0.12）。 对于小型模型（具有单个子网），从模型在移动设备上的部署角度来看，这是有希望的。

消融研究表明了自组装技术的重要性。 平均而言，它改善了PSNR（0.21）和SSIM（0.0002）指标。 有趣的是，与没有NAS的基准相比，没有自组装的NAS模型的性能要差得多。

7.潜在的改进和未来的工作

尽管我们的搜索过程显示出对小模型的最大PSNR改进，但这是以选择卷积内核为代价的，该卷积内核具有更大的内核大小和更多的滤波器。 但是，内核大小的可微表示可以将其他计算约束引入优化过程。 我们期望配备额外计算正则化的训练将导致找到可在便携式设备上部署的模型。

我们的搜索过程类似于贝叶斯深度学习的变分推理方法； 因此，我们的方法可能会遭受[12]中报道的对搜索空间的探索不足。 当我们的方法仅探索接近梯度下降轨迹的点时，可能会出现此问题，这会缩小搜索空间的探索范围，并使其可能对权重初始化敏感。 我们计划在将来的工作中探索这种敏感性。 此外，在我们的训练过程中，权重和结构参数都得到了优化。 我们计划探索两级优化（类似于[22]），以检查这些参数的耦合如何影响结果。

由于不同超级内核中的采样是独立执行的，因此可能会出现与公平有关的问题[8]。 因此，根据马修定律，网络可能会迅速收敛到随机的次优解决方案。 我们计划测试更高级的概率模型（例如基于马尔可夫随机场，递归神经网络或注意力），以便为更高级的搜索空间建模并进行更广泛的探索。

此外，我们观察到培训过程可能会对最终结果产生重大影响。 具体而言，与非超级内核模型相比，配备超级内核的模型可能需要更长的训练时间（更高的耐心）。 我们怀疑这可能是导致NAS型号性能下降而没有自我组装的原因。 我们计划在以后的工作中对此主题进行调查。

8.结论

在这项工作中，我们引入了一种快速轻量级的神经体系结构搜索图像去噪算法。我们证明，它可以取得远远超过无NAS解决方案的最新结果。 我们已经证明，提出的超级内核技术可以在几个GPU小时内达到与最新架构的图像去噪相当的结果。