NAS-DIP: Learning Deep Image Prior with Neural Architecture Search

NAS-DIP: 用神经结构搜索学习深度图像先验

论文链接：https://arxiv.org/abs/2008.11713

项目链接：https://github.com/YunChunChen/NAS-DIP-pytorch

Abstract

最近的研究表明，深度卷积神经网络的结构可以用作解决各种逆图像恢复任务的结构化图像先验。我们建议寻找能够捕获更强图像先验的神经结构，而不是使用手工设计的结构。在通用U-Net架构的基础上，我们的核心贡献在于为(1)上采样单元和(2)跨尺度残差连接模式设计新的搜索空间。我们通过利用现有的神经结构搜索算法(使用带有循环神经网络控制器的强化学习)来搜索改进的网络。我们通过各种各样的应用验证了我们方法的有效性，包括图像恢复、去雾、图像到图像的翻译和矩阵分解。大量的实验结果表明，我们的算法优于目前最先进的自由学习方法，在某些情况下可以与现有的基于学习的方法相媲美。

1 Introduction

卷积神经网络(CNN)已经成功地应用于许多计算机视觉任务。除了图像识别任务，CNN在图像恢复和合成问题上也表现出了很强的性能。这些成功案例背后的原因通常归因于CNN从大规模数据集中学习先验的能力(即，先验被嵌入到训练网络的参数/权重中)。与现有的需要从标记数据集中学习网络参数的监督学习算法不同，最近的研究发现，网络的结构本身就足以捕获丰富的低级图像统计[67,75]。这种在网络架构中编码的结构化图像先验对于图像恢复(例如，单幅图像超分辨率[44,45,46]和图像去噪[12,49,79])和图像合成(例如，图像到图像的翻译[37,38,47,48,88])任务至关重要。

虽然与基于学习的方法相比，自由学习方法[75]在图像恢复任务上表现出了具有竞争力的性能[44,46]，但只有传统的网络架构，如ResNet[31]或U-Net[66]得到了评估。针对这些图像恢复问题，网络设计有两个重要方面。首先，虽然编码器的设计已经被广泛研究[31,35,43,69]，但解码器的设计[77, 78] (特别是上采样单元)受到的关注相当少。其次，随着特征的空间分辨率沿着特征编码器的路径逐渐降低，网络恢复具有更高空间分辨率的特征映射至关重要。U-Net[66]是一种解决这一问题的流行设计，它在执行一系列上卷积时，将相应编码器层的编码特征与解码器中的特征连接起来。然而，这种跳过连接模式是为手头的每个任务手工设计和固定的。

我们的工作。在本文中，我们建议搜索(1)解码器中的上采样单元和(2)编码器和解码器之间的跳跃连接模式(即跨级特征连接)。为了实现这一点，我们为这两个组件开发了新的搜索空间。首先，为了搜索上采样单元，我们将典型的上采样操作分解为两个步骤：i)改变空间分辨率的方法(如双线性、双三次上采样[20]或最近邻上采样)和 ii)特征变换的方法(如二维卷积或二维转置卷积[21,54,83])。其次，为了搜索跨层连接模式，我们提出在编码器-解码器网络中搜索跨不同特征层共享的连接模式。

神经架构搜索(NAS)算法[24,26,89]已被证明在大型搜索空间中发现具有最佳性能的网络是有效的，我们利用带有循环神经网络(RNN)控制器的强化学习(RL)[24,26,89]，并使用PSNR作为奖励来指导架构搜索。通过在两个开发的搜索空间中同时搜索，我们的方法能够发现一个CNN架构，该架构可以为感兴趣的任务捕获更强的结构化图像先验。我们通过四个无需学习的图像恢复任务，包括单幅图像超分辨率、图像去噪、图像补全和图像去雾，以及一个基于学习的未配对图像到图像的转换问题(见图1)，展示了我们方法的适用性。我们的实验结果表明，搜索上采样单元和跨层特征连接，比传统的神经结构性能有所提高。

我们的贡献。首先，我们提出了基于几种常用上采样算子的分解，使我们能够为每个任务搜索新的上采样单元。其次，我们开发了一个由编码器-解码器架构中的跨层特征连接模式组成的搜索空间。第三，对各种图像恢复和合成任务的广泛评估表明，我们提出的算法在某些情况下优于现有的基于学习的方法，并且与现有的自由学习方法相比达到了最先进的性能。

2 Related Work

上采样Cell。上采样单元的设计可分为两类：1)基于不可学习参数的方法和2)基于可学习参数的方法。基于不可学习参数的方法使用插值将特征图的大小从较低的空间分辨率调整到较高的空间分辨率。示例算子包括双线性/双三次插值[20,59]、最近邻上采样[9,39,59]和深度到空间上采样[46]。基于可学习参数的方法学习低空间分辨率和高空间分辨率特征映射之间的映射。在设计选择中，2D转置卷积是各种密集预测任务的热门选择之一，包括语义分割[54]、光流估计[21]、深度预测[15]、图像恢复[49,73]和合成[59]问题。最近的进展包括双线性加性上采样[77,78]和CARAFE[76]。

与这些手动设计上采样操作的方法相比，我们通过将几种现有的上采样操作解耦为改变空间分辨率的方法和特征变换的方法来开发上采样单元的搜索空间，并采用神经结构搜索算法[24,26,89] (即带有RNN控制器的强化学习)来自动发现每个单独任务的最优上采样单元。我们进一步证明，与基本网络架构相比，发现的上采样单元可以跨任务传输，并且具有良好的性能。

NAS应用和搜索空间。NAS算法已经成功地应用于各种任务，包括图像分类[53,61,63,64,89,90]、语义分割[14,51,58]、目标检测[24]、图像恢复[17,72]和图像生成[26]。在目标检测的背景下，NAS-FPN[24]开发了一个搜索空间，允许模型通过合并跨尺度特征来学习金字塔表示，以改进对不同尺度和位置的多个物体的检测。在语义分割的背景下，有几种方法侧重于搜索编码器架构[51]、Atrous空间金字塔池模块[14]或用于紧凑架构的解码器单元[58]。在图像恢复任务中，现有算法的目标是为单幅图像超分辨率寻找更好的编码器结构[17]，或为图像补全和图像去噪寻找非对称编码器-解码器结构[72]。在图像生成问题中，AutoGAN[26]采用渐进式方案来寻找更好的生成对抗网络(generative adversarial network, GAN)架构[27]。搜索解码器架构的其他方法侧重于搜索更紧凑的架构[58]或使用手工设计的上采样单元和特征连接优化单元结构[26，72]。

我们的侧重点与这些方法有两点不同。首先，我们为上采样单元开发了一个搜索空间，允许我们为手头的每个任务发现最优上采样选择。其次，我们寻找跨层特征连接的模式，并在编码器-解码器架构中的不同特征层之间共享它。我们的跨层特征连接与NAS-FPN[24]的不同之处在于，我们的目标是在解码器中恢复具有更高空间分辨率的特征映射，而NAS-FPN[24]的目标是学习用于目标检测的金字塔状特征表示。通过同时搜索上采样单元和跨层特征连接，我们的搜索架构在各种图像恢复任务上与现有的自由学习方法相比达到了最先进的性能，并且与现有的基于学习的算法相比也达到了具有竞争力的结果。

NAS算法。NAS方法可以根据搜索算法分为几类。主要的方法有进化[5,52,56,63,64,70,71,80]，强化学习[7,13,74,87,89,90]和可微搜索[2,53]。进化搜索利用进化算法，通过从一群体系结构中随机变异高性能候选者来发现网络结构。基于RL的方法采用Q-Learning[7,87]或策略梯度[13,74,89,90]策略来训练一个循环网络，该网络输出一系列描述网络结构的符号[89]或可重复的单元结构[90]。可微搜索方法开发了一个连续的搜索空间，并使用梯度下降优化网络架构，以消耗大量内存为代价提供了时间效率。

在这项工作中，我们遵循基于RNN的方法，采用基于RNN的搜索算法和RNN控制器[26,89]来搜索上采样单元和跨层连接模式。我们注意到搜索算法并不局限于基于强化学习的方法。也可以应用其他替代方法，如基于进化的方法或可微分体系结构搜索算法。本文的重点在于两个搜索空间的设计。我们把搜索算法的开发留给未来的工作。

3 Proposed Method

在本节中，我们首先概述所建议的方法。然后，我们分别描述了上采样单元和跨尺度残差连接的两个开发的搜索空间。

3.1 方法概述

与现有的直接从大规模数据集学习的基于学习的方法[44] (图2a)相反，最近的研究[75]表明，通过将噪声随机映射到退化(例如，有噪声、低分辨率或遮挡)的图像，未经训练的CNN可以解决具有竞争力的图像恢复问题(图2b)。为了发现捕获更强图像先验的网络结构，我们考虑在不从成对数据中学习的情况下搜索上采样单元和跨尺度残差连接模式的任务。为了实现这一点，我们提出了一个由两个步骤组成的算法。如图2c所示，我们首先使用RNN控制器[89] (使用PSNR作为奖励) 应用强化学习，在一个固定训练集(蓝色块)上搜索性能最佳的网络结构$f_{\theta }^{\ast }$。在网络结构搜索步骤之后，对于测试集中的每个图像，我们随机重新初始化性能最佳的网络结构$f_{\theta }^{\ast }$的权重，并优化从随机噪声到退化图像(绿色块)的映射。

搜索性能最好的网络结构。给定训练集中的图像$x\in {\mathbb{R}}^{H\times W\times 3}$，我们首先通过添加噪声、下采样或根据感兴趣的任务从$x$中删除某些像素来生成退化版本$x_0$。图像去噪，$x_0\in {\mathbb{R}}^{H\times W\times 3}$表示$x$的嘈杂的版本，针对单一图像超分辨率，$x_0\in {\mathbb{R}}^{\frac{H}{r}\times \frac{W}{r}\times 3}$表示低分辨率版本的$x$ ，$r$代表将采样率，对图像修复，$x_0\in {\mathbb{R}}^{H\times W\times 3}$表示$x$的遮盖版本。然后我们样本噪声图像$z\in {\mathbb{R}}^{W\times H\times C}$和执行搜索网络$f_{\theta }$绘制将噪点图像$z$映射到去噪后的、高分辨率的或补全版本的$x_0$，也就是说，将噪声图像$z$映射到$x$。为了实现这一点，我们遵循DIP[75]，并针对不同的任务优化不同的目标。

由于训练集中的真实图像是可用的，我们可以通过计算真实图像与网络输出(即$f_{\theta }(z)$)之间的峰值信噪比(PSNR)来对每个搜索到的网络结构进行排序，并确定训练集中表现最佳的网络结构$f_{\theta }^{\ast }$。

确定最优停止点$t^{\ast }$。我们注意到最优停止点$t^{\ast }$(即所需的迭代次数)取决于网络结构。由于我们有一个固定的训练集，我们能够估计每个随机生成的网络结构的最佳停止点。然后，我们通过测量恢复图像与其相应的真值(即来自训练集的原始图像)之间的差异(即计算PSNR)对所有采样的网络结构进行排序。之后，我们将性能最好的网络结构应用于测试集，并报告在最优停止点$t^{\ast }$记录的结果。

用搜索到的网络结构$f_{\theta }^{\ast }$进行测试。在训练集上搜索后，我们一次对测试集中的每个图像使用随机初始化的最佳网络结构$f_{\theta }^{\ast }$进行优化，使用与DIP[75]相同的目标对不同的任务进行$t^{\ast }$迭代。

3.2 上采样层的搜索空间

通过分解现有的上采样操作，我们建立了上采样层的搜索空间：1)改变特征图空间分辨率的方法和2)特征变换的方法。我们对改变空间分辨率的操作的搜索空间包括:双线性上采样、双三次上采样[20]、最近邻插值、深度到空间[46,68]和步幅2转置卷积[21,54,83]。我们的特征变换操作搜索空间包括：二维卷积、每N个连续通道相加[77,78]、可分卷积[28,77,78]、深度卷积[23,28,29]、恒等。为了在网络架构搜索过程中放松自由度，我们的搜索空间允许包含可学习参数的操作来搜索核大小、空洞率以及是否在最后包含激活。激活函数的操作搜索空间包括None、ReLU[57]、LeakyReLU[81]、SELU[42]和PReLU[30]。图3展示了我们为上采样单元开发的搜索空间。通过分解几种常用的上采样算子，我们的搜索空间更加灵活，并允许我们为每个任务发现一个新的上采样单元。新开发的空间上采样算子(例如CARAFE[76])也可以在未来很容易地纳入我们的搜索空间。

3.3 跨尺度残差连接

对于跨尺度残差连接，我们开发了一个覆盖跨尺度特征连接的搜索空间。与U-Net[66]在同一特征级别连接特征映射不同，我们采用残差连接[31]。这允许我们合并从不同特征层提取的特征映射，而不需要预先定义解码器中每层的输入通道数量，因为我们设计的输入通道数量总是相同的。

共享跨层模式。跨尺度残差连接的搜索空间可能非常大。假设网络深度为d(即编码器中有$d$个特征层，解码器中有$d$个特征层)，则可能的连接模式数为$2^{d^2}$。为了限制搜索空间，我们强制连接模式仅依赖于特征级别的差异(例如，将所有特征映射连接到更高一层的特征映射)。对于解码器中的第$k$个特征层，它可以接收从$(k-d+1)$特征层到$(k+d-1)$特征层的特征映射(即总共有$2^{2d-1}$种可能的连接模式)。有了这个约束，我们就可以搜索特征连接的模式，并在解码器的不同特征层之间共享这个模式。图4说明了建议的跨层残差连接的主要思想。搜索空间的大小可以从$2^{d^2}$(没有任何连接模式的限制)显著减少到$2^{2d-1}$。

渐进式上采样。对于跨尺度上采样操作，如图5所示，我们建议将左侧的4x上采样操作解耦为中间的两个连续的2x上采样操作，并在相同的特征级别上与右侧的2x上采样操作共享权重。这允许我们只在每个连续的特征级别之间定义2倍上采样操作，所有其他可能的上采样尺度都可以通过将它们解耦成一系列的2倍上采样操作来实现。跨尺度下采样连接也可以类似地实现。

4 Experimental Results

在本节中，我们首先描述实现细节。然后，我们提出了定量和视觉比较现有的方法以及消融研究。

4.1 实现细节

在这里，我们提供了关于神经结构搜索、模型训练和测试的实现细节。

神经结构搜索。我们使用PyTorch实现我们的模型。我们的RNN控制器的网络架构与[26]相同。为了创建用于搜索网络架构的训练数据，我们从DIV2K[1]训练集中随机抽取100张图像。为了寻找目标任务的最优网络结构，神经结构搜索过程由两个交替的阶段组成。对于第一阶段，RNN控制器首先对随机初始化的候选网络体系结构进行采样。然后，我们在hold-out训练集(即我们的NAS-DIP框架中的模型训练)上优化采样的候选模型。对于第二阶段，我们首先计算恢复的预测和相应的基础真值之间的PSNR作为奖励，并使用强化学习来更新RNN控制器(即我们的NAS-DIP框架中的NAS训练)。每个任务所需的训练时间各不相同。具体来说，找到超分辨率任务的最佳网络结构大约需要3天，去噪大约需要3天，补全大约需要5天(使用内存为12GB的NVIDIA V100 GPU)。

测试的细节。当应用搜索到的模型进行测试时，可以设置一个超参数(迭代次数)来获得最终的预测结果。早期的方法[75]利用测试图像的ground truth和PSNR来选择性能最好的迭代次数。然而，这种方案可能不实用，因为ground truth图像并不经常可用。为了解决这个问题，我们使用相同的训练集进行NAS训练，以找到允许模型达到最佳性能的最佳迭代次数。具体来说，我们选择模型预测值为超分辨率4500次迭代，去噪3500次，9000次补全迭代。

4.2 定量比较

我们通过评估其在解决图像恢复中的各种逆问题之前用作深度图像时的性能来验证我们搜索的模型架构的有效性。在每个任务中，我们在基准数据集上与最先进的自由学习方法进行比较。

单图像超分辨率。根据DIP[75]中的评估协议，我们采用了两个标准基准：Set5[10]和Set14[84]数据集。我们将我们的方法与现有的自由学习方法[8、20、25、36、55、65、75]和基于学习的方法[19、41、44、50、62、85、86]在三种不同的上采样尺度(即2倍、4倍和8倍上采样)上进行了比较。

我们使用DIP[75]提供的评估代码。表1a给出了实验结果。在所有三个上采样尺度上的结果表明，我们提出的算法与最先进的自由学习方法相比表现良好。我们的研究结果表明，对于更大的上采样规模(例如8倍上采样)，与现有的基于学习的方法相比，我们的方法在两个数据集上都取得了竞争甚至更好的性能[41,44]。这很有趣，因为我们的模型从未见过任何配对的低/高分辨率图像对。结果还突出了我们的搜索架构的重要性，与使用手工制作先验的现有方法(例如DIP[75])相比，显著提高了性能。

图像去噪和补全。我们采用BM3D数据集[18]来评估图像去噪任务为了公平比较，我们遵循DIP[75]，并使用指数滑动窗口(移动平均)平均我们的网络输出以获得最终结果。为了评估补全任务的性能，我们遵循DIP[75]并使用Heide等人[33]提供的数据集在这里，我们将我们的结果与DIP [75]， Papyan等[60]和SGLD[16]使用50%缺失像素设置进行比较。表1b给出了实验结果。同样，我们的方法在这两个任务上都优于所有竞争方法。

与最近CNN设计的对比。最近有几种方法设计CNN架构以提高图像恢复任务的性能。我们首先遵循与DeepDecoder[32]相同的实验设置，并在4倍超分辨率和补全任务上评估我们的方法表1c报告了实验结果。接下来，我们遵循潜卷积模型[6]中的评估方案，并报告表1d中4倍超分辨率和补全任务的结果。

从广泛的定量评估中，我们表明，我们的模型具有搜索到的上采样单元和跨尺度残差连接，可以在现有的手动CNN架构设计之前作为更强的结构化图像。

4.3 视觉对比

在这里，我们展示了几个图像恢复任务的样本定性结果，并将它们与最先进的方法进行比较。我们建议读者回顾全分辨率结果，以更好地感知视觉质量的改善。

图6和图7分别是单幅图像超分辨率和图像去噪的视觉结果。一般来说，使用我们的模型作为图像先验，在视觉质量方面会有明显的改善。这一改进突出了通过神经结构搜索学习更强的结构化图像的优势。

除了标准的图像恢复任务，我们还实验了模型在两个不同任务中的可转移性。我们使用DoubleDIP[22]中的去雾应用和CompMirror[3]中的矩阵分解任务进行演示。

对于去雾任务，我们遵循DoubleDIP[22]的官方实现来生成去雾结果为了生成我们的结果，我们将DoubleDIP[22]中的标准U-Net模型与我们在去噪任务中搜索的搜索模型交换。图8显示了在O-HAZE数据集[4]上与DoubleDIP[22]进行视觉比较的示例。结果表明，使用该模型产生的去雾图像具有更好的视觉质量。

对于矩阵分解，我们使用CompMirror[3]的官方实现来生成分解的结果为了生成我们的结果，我们将CompMirror[3]模型中的上采样层替换为我们在去噪任务中搜索到的上采样层。图9给出了一个与CompMirror[3]进行可视化比较的例子。我们的结果表明，使用我们的模型产生更平滑的分解图像，视觉伪影更少。

4.4 消融研究

我们进行消融研究，以分离来自各个组成部分的贡献。具体来说，我们的目标是了解我们的两项技术贡献在多大程度上提高了性能。由于我们的方法建立在DIP的U-Net架构上[75]，我们使用他们的结果作为基线。

我们在表2中报告了不同方法的结果。我们的研究结果表明，搜索上采样单元和跨层残差连接模式始终有助于提高跨多个任务的DIP性能[75]。我们还观察到上采样单元对于单幅图像的超分辨率任务尤为重要。另一方面，引入搜索的跨尺度残差连接，对于上采样单元的涂漆和去噪任务都提供了更大的性能提升。具有两个组件的模型显示出最佳性能，突出了这两个组件之间的互补性。

4.5 图像到图像的翻译

我们还探讨了将搜索到的模型(从去噪任务)转移到另一个问题。具体来说，我们的目标是测试我们的搜索模型是否可以很好地推广到图像到图像的翻译任务。我们采用作者提供的CycleGAN[88]的PyTorch实现，并训练了Summer <-> Winter翻译我们将我们的结果与CycleGAN的标准U-Net模型进行了比较[88]。图10显示了每个翻译方向的一个示例结果。

为了量化性能，我们还计算了FID评分[34]并进行了用户研究。我们在表3中报告了结果。客观(FID评分)和主观(用户研究)结果都表明，我们搜索的模型比基本CycleGAN[88]模型的性能有所提高。

5 Conclusions

在本文中，我们提出使用神经结构搜索技术来发现由CNN结构捕获的更强的结构化图像先验。我们工作的核心技术贡献在于：(1)解码器中常用的上采样层的搜索空间设计；(2)编码器和解码器之间的跨级特征连接。我们在标准的U-Net架构上构建我们的网络设计，并为每个感兴趣的任务寻找最佳的上采样单元和跨层特征连接模式。我们在四个图像恢复任务、一个矩阵分解应用和一个不配对的图像到图像的翻译问题上验证了我们的模型的有效性。通过广泛的实验评估，我们的结果表明，与传统的网络设计相比，我们的性能得到了持续的改善。