[论文解读] Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning

Ke Yu, Chao Dong+2 more|arXiv (Cornell University)|Apr 10, 2018

Advanced Image Processing Techniques参考文献 41被引用 31

一句话总结

本文提出 RL-Restore，一种强化学习框架，可从一组轻量级、任务特定的卷积神经网络工具箱中动态选择，以恢复具有复杂混合失真的图像。通过将恢复过程建模为基于逐帧 PSNR 奖励的序列决策过程，并联合训练智能体与工具，该方法在仅使用大 CNN 82.2% 计算量的情况下实现了相当的性能，同时提升了对未见伪影的鲁棒性。

ABSTRACT

We investigate a novel approach for image restoration by reinforcement learning. Unlike existing studies that mostly train a single large network for a specialized task, we prepare a toolbox consisting of small-scale convolutional networks of different complexities and specialized in different tasks. Our method, RL-Restore, then learns a policy to select appropriate tools from the toolbox to progressively restore the quality of a corrupted image. We formulate a step-wise reward function proportional to how well the image is restored at each step to learn the action policy. We also devise a joint learning scheme to train the agent and tools for better performance in handling uncertainty. In comparison to conventional human-designed networks, RL-Restore is capable of restoring images corrupted with complex and unknown distortions in a more parameter-efficient manner using the dynamically formed toolchain.

研究动机与目标

• 解决单个大型 CNN 在处理复杂混合图像失真时存在的参数量和计算成本过高的局限性。
• 通过动态组合一系列轻量级、任务特定的网络，实现自适应且参数高效的恢复。
• 通过智能体与工具的联合训练，提升对未知或未见失真在恢复过程中的鲁棒性。
• 通过可见的工具链选择过程，揭示逐步恢复流程，提升可解释性。

提出的方法

• 该框架使用一组小型专用 CNN 工具，分别针对去模糊、去噪和 JPEG 伪影减少任务进行训练。
• 智能体学习一种策略，根据当前图像状态和先前动作，逐步选择工具，使用与 PSNR 提升成比例的逐帧奖励函数。
• 采用联合学习方案，同时训练智能体与工具，以更好地应对中间恢复状态中的不确定性和伪影。
• 当图像质量被认为足够好时，智能体可自主停止过程，防止过拟合并节省计算量。
• 奖励函数通过 PSNR、SSIM、最终 PSNR 和最终 MSE 进行评估；逐帧 PSNR 表现最佳。
• 工具的训练数据包含合成噪声和压缩数据，以增强对中间伪影的鲁棒性。

实验结果

研究问题

• RQ1 强化学习智能体能否有效组合轻量级网络的动态工具链，以恢复具有复杂混合失真的图像？
• RQ2 智能体与工具的联合训练如何提升对未知或中间伪影的鲁棒性？
• RQ3 不同奖励函数对恢复性能和收敛性有何影响？
• RQ4 自动停止机制如何影响恢复质量与计算效率？
• RQ5 该框架在泛化到未见失真方面是否优于单个大型 CNN？

主要发现

• RL-Restore 在仅使用大 CNN 82.2% 的 FLOPs 情况下，实现了与之相当的 PSNR 性能，证明了显著的计算效率。
• 联合训练策略使 PSNR 相比基线工具提升 0.2 dB，若在训练数据中加入噪声增强，还可额外提升 0.2 dB。
• 逐帧 PSNR 奖励函数在 PSNR 和 SSIM 指标上均优于逐帧 SSIM、最终 PSNR 和最终 MSE。
• 若移除自动停止机制，PSNR 平均降低 0.15 dB，且在轻微失真情况下下降更明显，表明存在过恢复风险。
• 长度为 12 的工具链搭配 3 种工具可实现最优性能，更长的链因训练复杂度增加而未显著提升结果。
• 该框架揭示了可解释的、逐步的恢复过程，相较于黑箱 CNN 提供了更高的透明度。