Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning

摘要

        我们通过强化学习研究了一种新的图像恢复方法。 与现有的大多数为特定任务训练单个大型网络的研究不同，我们准备了一个由不同复杂性的小规模卷积网络组成的工具箱，专门负责不同的任务。 我们的方法RL-Restore然后学习一个策略，从工具箱中选择适当的工具来逐步恢复损坏图像的质量。 我们制定逐步奖励函数，与每一步恢复图像的程度成正比，以了解行动政策。 我们还设计了一个联合学习计划来训练代理和工具，以便在处理不确定性方面获得更好的性能。 与传统的人工设计网络相比，RL-Restore能够使用动态形成的工具链以更高参数效率的方式恢复被复杂和未知失真破坏的图像。

1、介绍

        深度卷积神经网络（CNN）取得了巨大的成功，不仅在高级视觉任务中，而且在低级视觉任务（如去模糊、去噪、JPEG伪像减少和超分辨率方面）。 特别是，与传统的基于模型的优化方法相比，证明了良好的性能和快速的测试速度。

        由于CNN的辨别性，大多数这些模型都经过训练以处理专门的低级视觉任务。 例如，在JPEG伪像减少中，已经设计了用于不同压缩质量的不同网络以实现令人满意的恢复。 在超分辨率的情况下，通常具有不同的网络来处理不同的缩放因子。 最近的一些研究表明，使用CNN可以同时处理多种失真类型或应对不同程度的降级。 然而，这通常伴随着使用更深层网络的消耗。 此外，这样的网络处理具有相同结构的所有图像，尽管其中一些图像较不困难并且可以以花费较少的方式恢复。

        在本文中，我们探讨了使用一些较小规模但专门的CNN来协同解决更难恢复任务的可能性。 我们的想法背离了当前的理念，即需要大容量CNN来解决复杂的恢复任务。 相反，我们希望拥有一套工具（基于小型CNN），并学会自适应地使用它们来解决手头的任务。 上述想法可以提供新的见解，CNN如何用于解决现实世界的恢复任务，其中图像可能被混合的失真污染，例如，在几个处理阶段之后产生的噪声和块效应。 此外，与现有的基于CNN的模型相比，新方法可以导致参数有效的恢复。 特别是，可以根据失真的严重程度选择不同复杂度的工具。

        为实现这一目标，我们提出了一个框架，将图像恢复视为一个决策过程，通过该框架，代理可以自适应地选择一系列工具来逐步细化图像，如果恢复的质量被认为是令人满意的，代理可以选择停止。 在我们的框架中，我们准备了许多具有不同复杂性的轻量级CNN。 它们是特定于任务的，旨在处理不同类型的恢复分配，包括去模糊，去噪或JPEG伪像减少。 选择工具的顺序是在强化学习（RL）框架中制定的。 代理通过分析当前步骤中恢复的图像的内容并观察所选择的最后一个动作来学习决定下一个要选择的最佳工具。 当代理改善输入图像的质量时，累积奖励。

        我们将提议的框架称为RL-Restore。 我们总结了我们的贡献如下：
        1）我们提出了在强化学习框架中解决图像恢复的新尝试。 与部署单个且可能较大的网络结构的现有方法不同，RL-Restore可以灵活地使用不同容量的工具来实现所需的恢复。
        2）我们提出了一个联合学习方案来同时训练代理和工具，以便框架具有更好的能力来应对在处理过程中出现的新的和未知的工件。
        3）我们表明，动态形成的工具链与强大的人工设计网络竞争性地具有较低的计算复杂度。 我们的方法可以在一定程度上应对看不见的失真。 有趣的是，我们的方法比现有方法更透明，因为它可以揭示如何使用不同的工具逐步消除复杂的失真。

        图1（b-c）说明了恢复被多次失真破坏的图像的学习策略，其中图像质量逐步完善。 两个基线CNN模型的结果如图1（d-e）所示，其中（d）具有与我们相似的参数数量（代理+工具应用），而（e）有两倍以上。 正如我们将在实验部分中进一步介绍的那样，RL-Restore优于CNN方法，给出相似的复杂性，并且它需要82.2％的计算来实现与单个大型CNN相同的性能。

2、相关工作

        CNN用于图像恢复：图像恢复是一个广泛研究的主题，旨在从损坏/嘈杂的观察中估计清晰/原始图像。 基于卷积神经网络（CNN）的方法已经在各种图像恢复任务中表现出优异的性能。 这些研究中的大多数训练专门研究手头任务的单个网络，例如去模糊，去噪，JPEG伪像减少和超分辨率。 我们的工作提供了一种更具参数效率但适应失真形式的替代方案。

        有几项开创性研究同时处理多种退化问题。 Kim等人通过开发20层深的CNN。 [19]使用单一模型来处理多尺度图像超分辨率。 [10]构建一对多网络，可以处理具有不同级别的压缩工件的图像。 [44]提出了一个20层深的CNN来同时解决多个恢复任务，包括图像去噪，JPEG伪像减少和超分辨率。 这些研究都没有考虑混合失真，其中单个图像受到多个失真的影响。 与上述工作不同，我们有兴趣探索是否可以使用3到8层的小规模CNN来联合恢复被混合失真污染的图像。

        存在可用于将大型网络压缩到较小网络以提高计算效率的方法。 在图像恢复领域，研究递归神经网络以减少网络参数。 然而，由于大量递归，计算成本仍然很高。 我们的工作目标与前面提到的研究正交，我们的框架通过学习策略来决定为任务选择合适的CNN而不是压缩现有的任务来保存参数和计算。

        深度强化学习： 强化学习是学习代理做出连续决策以最大化累积奖励的有力工具。 RL的早期工作主要集中在机器人控制上。 最近，传统的RL算法被纳入深度学习框架，并成功应用于各种领域，如游戏代理和神经网络架构设计。 还要注意计算机视觉领域的深度RL。 例如，Huang等人使用RL来学习加速CNN对象跟踪的早期决策策略。 Cao等人在低水平视觉中探索深度RL算法，并将注意机制应用于面部幻觉。 在本研究中，我们研究了RL框架中的恢复工具选择。 这个问题在文献中是新的。

3、学习修复工具链

        问题定义：给定图像I dis失真，我们的目标是恢复接近真实图像I gt的清晰图像I res. 失真过程可以表示为：

        其中◦表示函数组成，D 1，...，D n中的每一个表示特定类型的失真。与专注于单一类型失真的现有方法相比，我们打算处理多种失真的混合（即，n> 1）。例如，最终输出图像可能依次受到离焦模糊，曝光噪声和JPEG压缩的影响。在这种情况下，失真的数量n是3，并且D 1，D 2，D 3分别表示模糊，噪声和压缩。 为了解决混合失真问题，我们建议使用一系列恢复工具逐步恢复损坏的图像。

        挑战：工具选择的任务非常重要，并且对RL提出了独特的挑战。首先，恢复类型，级别和处理顺序的选择都会影响最终的性能。图2中显示了一个示例，其中图像被两种不同的失真组合破坏。使用适当的工具链，如图2（c，d）所示，图像质量和峰值信噪比（PSNR）值依次得到改善。然后我们稍微重新排列工具顺序，如图2（b，e）所示，或者调整工具的恢复级别，如图2（a，f）所示。结果表明，工具链中的微小变化会严重影响修复性能。具体而言，使用不正确的工具可能会导致不自然的输出，例如图2（a）中的过度锐化和图2（f）中的模糊。即使工具选择得当，不恰当的顺序也会降低性能（图2（b，e））。由于工具链的顺序会极大地影响结果，因此选择在每个步骤中使用哪种工具变得至关重要。

        当工具在特定任务上进行训练时，我们遇到另一个问题，即没有任何工具可以完美地处理“中间状态”，这是指经过几个处理步骤后的中间结果。 由于大多数失真是不可逆转的，因此图片的恢复并不是相应修复体的简单组合。可以在中间状态引入新的工件。例如，去模糊操作也会增强噪音，导致以下降噪器无法移除新引入的工件。 这里说到的是我们的任务所独有的。

        为了解决第一个挑战，我们将顺序工具选择问题视为马尔可夫决策过程（MDP），并以深度强化学习方式解决它。 为了解决第二个挑战，我们提出了一个培训方案，共同完善代理和工具，以便工具更好地了解代理可观察到的中间状态。 我们首先概述所提议的框架如下。

        RL-Restore概述： 建议的框架旨在发现给定输入图像损坏的工具链。 如图3所示，RL-Restore由两部分组成：1）工具箱（Toolbox），包含各种用于图像恢复的工具; 2）具有循环结构的代理（Agent），可在每一步或早期停止动作中动态选择工具。 我们将工具选择过程作为强化学习过程 - 对工具选择进行一系列决定，以最大化与恢复图像的质量成比例的奖励。 接下来，我们首先描述工具箱的合理设置，然后解释代理的详细信息。

3.1工具箱（Toolbox）

        工具箱包含一组可能应用于损坏图像的工具。 我们的目标是设计一个功能强大且轻量级的工具箱，因此我们限制每个工具精通特定任务。 也就是说，每个工具仅在很小范围的失真上进行训练。 为了进一步降低整体复杂性，我们使用更小的网络来实现更轻松的任务。为了我们的研究目的，我们准备了12个工具，如表1所示，其中每个工具被分配用于处理特定级别的高斯模糊，高斯噪声或JPEG压缩。 我们应用三层CNN（如[8]中所述）用于轻微失真，并应用更深的八层CNN用于严重失真。 注意，不必限制工具来解决上述失真。 我们做了这些选择，因为它们在图像恢复的文献中通常被认为是典型的。 在实践中，可以根据手头的任务设计具有适当复杂性的工具。

        如第二节第三段开头所讨论的那样。一组有限的工具并不能很好地处理中间状态下出现的新工件。 为了解决这个问题，我们提出了两种策略：1）为了提高工具的鲁棒性，我们在所有训练数据中添加了轻微的高斯噪声和JPEG压缩；2）训练代理后，所有工具都在基础上进行联合微调。 训练有素的工具链。 然后，这些工具将更加适应代理任务，并能够更加健壮地处理中间状态。 我们将在3.3小节中讨论训练步骤， 在第四章验证拟议策略的有效性。

3.2代理（Agent）

        RL-Restore的处理流程如图3所示。给定输入图像，代理首先从工具箱中选择一个工具并使用它来恢复图像，然后代理根据之前的结果选择另一个工具并重复恢复过程直到它决定停止。 我们将首先说明一些术语，如行动（Action），状态（State）和奖励（Reward ），然后详细介绍代理结构和恢复程序。

        行动（Action）：表示为A的动作空间是代理可以采取的所有可能动作的集合。 在每个步骤t，选择动作a t并将其应用于当前输入图像。 每个操作代表工具箱中的工具，还有一个表示停止的附加操作。 如果工具箱中有N个工具，则A的基数为N + 1.因此，代理的输出v t是一个（N + 1）维向量，表示每个动作的值。 一旦选择了停止动作，将终止恢复过程，并且当前输入图像将成为最终结果。

        状态（State）：状态包含代理可以观察到的信息。 在我们的公式中，状态被表示为S t = {I t，~v t}，其中I t是当前输入图像，并且~v t是过去的历史行为向量。 在步骤1，I 1是输入图像，并且~v 1是零向量。状态为代理人提供丰富的背景知识。 1）当前输入图像I t是必要的，因为所选动作将直接应用于该图像以获得更好的恢复结果。 2）前一动作矢量~v t的信息，即t-1步的代理的输出值矢量，即~v t = v t-1，也是重要的。 对先前决定的了解可以帮助在当前步骤中选择动作。 发现这比仅使用I t更有经验。

        奖励（Reward ）：奖励驱动代理的培训，因为它学会最大化累积奖励。 代理应该学习一个好的方式，以便最终恢复的图像令人满意。 我们希望确保每一步都提高图像质量，因此逐步奖励的设计如下：

        其中r t是步骤t的奖励函数，P t +1表示在第t步恢复结束时I t+1与参考图像I gt之间的PSNR，P t表示步骤t的输入PSNR。 累积奖励可以写为，这是恢复过程中的总体PSNR增益，并且它被最大化以实现最佳增强。 请注意，使用其他图像质量指标（例如，感知损失，GAN损失）作为我们框架中的奖励是灵活的。 调查超出了本文的重点。

        结构体（Structure）：在每个步骤t，代理商在给定输入状态S t的情况下评估每个动作的值，其可以表述如下：

        其中f ag表示代理网络，W ag表示其参数。 向量v t表示动作的值。 具有最大值的动作被选择为at，即，其中v t，a表示对应于动作a的值向量v t的元素。

        代理由三个模块组成，如图3所示。第一个模块，称为特征提取器，是一个四层CNN，后面是一个输出32维特征的完全连接（fc）层。 第二个模块是具有N + 1维输入和N维输出的单热编码器，保留先前所选动作的信息。 请注意，输出是一个低于输入的维度，因为在前一步骤中不能采用停止操作，因此我们只是删除最后一个维度。 前两个模块的输出连接到第三个模块的输入端，这是一个长短期存储器（LSTM）[15]。 LSTM不仅观察输入状态，还将历史状态存储在其存储器中，其提供历史恢复图像和动作的上下文信息。 最后，利用LSTM之后的另一个fc层，导出值向量v t以用于工具选择。

        恢复（Restoration）： 一旦根据v t中的最大值获得动作a t，相应的工具将应用于输入图像I t以获得新的恢复图像：

        其中f r表示恢复功能，W r表示工具箱中工具的参数。 如果选择了停止操作，则f r表示标识映射。 通过分别表示I dis和I res作为输入失真图像和最终恢复输出，整个恢复过程可表示为：

        其中f = [f ag; f r]且W = [W ag; W r]。 T是选择停止动作的步骤。 我们还设置了最大步长T max以防止过度恢复。 当t = T max并且未选择停止动作时，我们将在当前步骤之后终止恢复过程。 换句话说，我们添加一个T≤Tmax的约束。

3.3训练

        工具的训练遵循[19]中的标准设置，其中均方误差（MSE）被最小化。真实有效的图像（ps：The ground truth image），输入图像和工具分别表示为y，x和h。至于代理，训练是通过深度Q学习来解决的，因为我们没有关于正确选择行为的先验知识。在提出的框架中，v t的每个元素都是[30]中定义的动作值，因此损失函数可以写成L =（y t - v t，a t）^ 2其中

        γ= 0.99是折扣因子。我们还使用目标网络来稳定训练，这是fag的克隆，并在训练时每C步更新其参数。在上式中，v t + 1，a'来自f'ag和v t，t来自f ag。在训练时，从重放记忆中随机选择剧集，[12]中提出了两种更新策略，其中“随机更新”是指从每集的随机点更新并进行固定数量的步骤，“顺序更新”表示所有更新都在剧集开始时开始并进入其终点。在[12]中，声称两种更新策略都具有相似的性能。由于我们的工具链不会太长，我们只需采用“顺序更新”，其中每个训练序列都包含一个完整的工具链。

        联合训练：正如3.1节所讨论的，没有一个工具可以完美地处理中间状态，在中间状态中可能会在之前的恢复步骤中引入新的和复杂的工件。为了解决这个问题，我们提出了一种联合训练算法，如算法1所示，以端到端的方式训练工具，以便所有工具都能学会处理中间状态。具体地，对于批处理中的每个工具链，转发失真图像I 1以获得恢复结果I T + 1。给定最终的MSE损失，然后渐变沿着相同的工具链向后传递。同时，每个工具的梯度在一批中累积，最后使用梯度的平均值来更新相应的工具。上述更新过程重复进行几次迭代。

          实施细节：在我们的实现中，工具的训练类似于[19]，其中所有实验运行超过80个时期（3.2×10 5次迭代），批量大小为64.初始学习率为0.1并且其减少0.1倍每20个时代。对于联合训练，我们在算法1中设置M = 64，α= 0.0001，分别表示批量大小和学习率。联合训练超过2×10 5次迭代。在训练代理时，我们使用Adam优化器和批量大小为32.最大步长T max根据经验设置为3，重放内存的大小选择为5×10 5。更新频率C = 2,500，以便每2,500次迭代从最新的代理网络f ag复制目标网络f'ag。在5×10 5次迭代中，学习速率从2.5×10 -4到2.5×10 -5呈指数衰减。

4、实验

        数据集和评估指标： 我们对DIV2K数据集进行了实验，这是最新的大规模高质量图像恢复数据集.800 DIV2K训练图像分为两部分：1）前750个训练图像和 2）其余50张图片供测试。 DIV2K验证图像用于验证。 通过缩小因子2,3和4来增强训练图像。然后将图像裁剪成63×63子图像，分别形成具有249,344和3,584个子图像的训练集和测试集。

        我们采用混合失真进行代理训练和测试。 具体地，将高斯模糊，高斯噪声和JPEG压缩的序列添加到具有随机级别的训练图像中。 高斯模糊和高斯噪声的标准偏差分别在[0,5]和[0,50]中均匀分布，而JPEG压缩的质量在[10,100]中受到均匀分布。 所有混合扭曲分为五组，如图4所示，从极端轻微到极端严重。 我们抛弃了两个极端情况，这些情况要么太容易，要么太难以恢复。 对温和组进行训练和测试。 为了进一步测试概括能力，我们还对未包含在训练数据中的轻度和重度组进行测试。

        对比：我们将RL-Restore与DnCNN和VDSR进行比较，它们是图像恢复和超分辨率的最先进模型，并且它们都能够处理多种降级。 DnCNN和VDSR在20个卷积层中具有相似的结构，而在DnCNN中采用批量归一化。他们的参数超过60万（见表2）。相比之下，RL-Restore（包括代理和所选工具2）的复杂性仅为DnCNN和VDSR的三分之一，总共有0.19百万个参数。当我们参考63×63输入图像上的乘法次数时，可以在计算上观察到更大的间隙。为了与RL-Restore进行公平比较，我们将VDSR从20层缩小到15层（每层42个过滤器），形成一个名为VDSR-s的新基线，其具有与RL-Restore类似的复杂性。遵循[19,44]中的相同培训策略，我们首先使用代理培训集训练基线。然后我们使用代理和工具训练集对模型进行微调，直到收敛。

4.1合成数据集的定量评估

        我们在表3中提供了RL-Restore和基线在不同测试集上的定量结果。轻度和中度集的结果表明我们的方法明显优于VDSR-s，与DnCNN和VDSR相当，证明了提出的RL-Restore 可以实现与深度CNN相同的性能，复杂度低得多。 值得注意的是，在严格测试集上，RL-Restore分别超过DnCNN和VDSR 0.2dB和0.3dB，其中在训练数据中未观察到失真。 这表明我们基于RL的方法在处理看不见的失真方面更灵活，而固定CNN更难以概括为看不见的情况。 视觉结果如图5所示。

        为了检测RL-Restore的内部行为，我们分析了每一步中工具选择的频率。 结果显示在图6中，其中x轴上的0-12表示表1中的12个工具，13表示停止动作。 从三个图表中可以看出，工具选择是多种多样的，并且所有工具都以不同的比例使用。 具体而言，在第一步中首选去模糊和去噪工具，而在第二步中经常选择去噪和去JPEG工具。 最后一步倾向于以大概率停止代理--47％。 有趣的是，在对看不见的数据进行测试时，最后一步的停止作用比例分别为轻度和重度测试集的60％和38％，这表明更严重和复杂的失真需要更长的工具链才能恢复。

4.2 现实世界图像的定性评价

        在实际情况中，图像总是被具有未知退化内核的各种复杂和混合噪声失真，使得当前方法的恢复任务极其困难。 所提出的基于RL的方法可以阐明可能的解决方案。 当现实世界的失真（例如，轻微的离焦模糊，曝光噪声和JPEG伪像）接近训练数据时，所提出的RL-Restore可以容易地推广到这些问题并且比单个CNN模型表现得更好。

        示例在图7中示出，其中输入图像与不同的失真（例如，模糊，噪声，压缩）组合，由智能电话捕获。 我们直接在这些真实世界的图像上应用训练有素的RL-Restore和VDSR，而无需对测试数据进行进一步微调。 很明显，我们的方法受益于灵活的工具链，对于恢复真实世界的图像更有效。 具体来说，图7（a，c）显示RL-Restore可以成功处理由曝光和压缩引起的严重伪影，而图7（b，d，e）表明我们的方法能够恢复模糊和复杂的混合噪声。 值得注意的是，当代理对恢复的质量有信心时，选择停止动作（图7（c，d，e））。我们认为，拟议的框架有可能通过更强大的恢复工具来处理更复杂的实际失真。

4.3消融研究（通过移除某个模型或者算法的某些特征，来观察这些特征对模型效果的影响）

        在本节中，我们将研究所提出的RL-Restore的不同设置，并提供有关超参数选择的一些见解。 为了更好地区分每个因素的有效性，我们排除了以下所有实验的联合培训策略。
        工具箱大小和工具链长度：工具箱的容量和恢复操作的数量决定了恢复性能。 我们可以选择改变工具链的长度和工具箱的大小。 如表4所示，在当前问题设置下，RL-Restore在N = 12和T max = 3时表现良好。 更少的工具和更短的工具链会降低性能。 更多工具和更长的工具链可实现相当的性能。 我们将这种现象归因于学习更复杂工具链的难度增加。值得指出的是，长度为2的工具链具有与温和测试装置上较长工具链相当的PSNR，表明轻微失真需要较少的恢复步骤。

        工具训练：如3.1节所述，我们提出了两种工具训练策略来消除中间状态下的复杂伪影：1）在工具训练数据中添加轻微的噪声和压缩。 2）与代理进行联合培训。 对照实验如表5所示，其中'Original'设置代表基线，'+Noise'采用第一种策略，'+Joint'使用它们。 很明显，在训练数据中添加噪声成功地将PSNR提高了0.2 dB，联合训练进一步推动了所有测试集的0.2 dB，证明了两种训练策略的有效性。

        奖励功能：我们通过实验发现，奖励函数的选择可以在很大程度上影响绩效。 除了提出的基于PSNR的逐步奖励外，我们还研究了其他奖励函数：1）逐步SSIM，其中奖励是每一步的SSIM收益; 2）最终PSNR，其中奖励是最后一步给出的最终PSNR收益; 3）如[4]中的最终MSE，其中奖励最终为负MSE。 我们自适应地调整学习率以获得不同的奖励。 从表6中可以看出，在PSNR度量上执行最差的逐步SSIM似乎不是奖励的好选择。 最终的MSE在PSNR上略胜一筹，但在SSIM上表现最差。 最终的PSNR实现了与提议的逐步PSNR奖励类似的性能。 尽管如此，我们并未声称PNSR是最佳奖励，并且还鼓励其他评估方法进行进一步比较。

        自动停止：停止操作使代理在对恢复结果有信心时可以灵活地终止恢复过程。由于这种灵活的停止机制，它可以防止图像过度恢复并节省大量计算。 为了证明其有效性，我们将结果与有/无停止动作进行比较。 从表7中可以看出，当去除停止动作时，PSNR值下降约0.15dB。 观察到温和测试集上的间隙大于其他测试集上的间隙。 这与我们的经验一致，即如果代理人没有及时停止，轻微的扭曲很容易被恢复。

5、结论

        我们提出了一种基于强化学习的图像恢复新方法。 与大多数现有的基于深度学习的方法不同，在我们的方法中，学习代理以动态地选择工具链以逐步恢复被复杂和混合失真破坏的图像。 在合成图像和真实图像上的广泛结果验证了所提出方法的有效性。 凭借其固有的灵活性，通过开发强大的工具和适当的奖励，所提出的框架可以应用于更具挑战性的恢复任务或其他低级视觉问题。