Review Learning : 推进一体化超高清图像恢复训练方法

Review Learning: Advancing All-in-One Ultra-High-Definition Image Restoration Training Method

 

摘要

一体化图像恢复任务变得越来越重要，特别是对于超高清（UHD）图像。

现有的一体机UHD图像恢复方法通常通过引入针对不同退化类型的即时或定制的动态网络来提高模型的性能。在推理阶段，它可能是友好的，但在训练阶段，由于模型在一个时代遇到多个不同质量的退化图像，这些杂乱的学习目标可能是模型的信息污染。

为了解决这个问题，本文提出了一种新的一般图像恢复模型的训练范式，本文将其命名为Review Learning，它使图像恢复模型能够在没有先验知识和提示的情况下处理多种类型的退化。该方法首先在几个退化数据集上对图像恢复模型进行顺序训练，并结合一种回顾机制，以增强图像恢复模型对以前几类退化数据集的记忆。此外，本文设计了一个轻量级的通用图像恢复网络，它可以有效地推理在单一消费级GPU上具有4K（3840×2160）分辨率的退化图像。

Introduction

近年来，先进的成像传感器和显示器的出现，极大地促进了超高清（UHD）成像技术的发展。然而，UHD图像中像素数量的增加无疑使它们在成像过程中更容易受到多个未知退化的影响。以前的UHD图像恢复方法通常分别处理这些基本的图像退化，包括低光图像增强、去雨、去雪和去模糊，通过使用特定的单任务模型。实际上，本文更需要一个一体化的模型来恢复不同类型的UHD退化图像。

最近的一系列工作专注于统一图像恢复，其中模型在混合退化数据集上进行训练，并对恢复过程中的退化类型进行隐式分类。虽然结果令人印象深刻，但批处理大小的减少以适应消费者级gpu的内存限制导致了训练时间的延长。此外，这种训练方法依赖于具有较强感知能力的提示信号，不幸的是，这些提示信号难以捕捉退化类型的图像。有鉴于此，本文对一体机图像恢复任务提出了两个关键问题：

i) 可以通过没有提示的混合退化图像来优化模型吗？

ii) 如何避免训练模型阶段的信息污染（多种退化类型的图像对模型记忆的干扰）？

为了解决这些问题，并受最近研究中证明的持续预训练（CPT）原则的启发，本文提出了一种名为Review Learning的方法，该方法可以逐步增强任意模型在统一图像恢复区域的能力，而不会产生显著的一般性能损失。具体来说，该模型学习一个特定的任务（如低光图像增强或图像去模糊），然后在这些知识的基础上学习新任务，定期回顾以前任务中具有挑战性的样本，以避免知识丢失。本文进行了全面的分析和统计，以确定审查困难样本的频率和间隔。这种训练模型的范式利用了神经网络的最大潜力，而不提示信息。

本文的贡献：

本文提出审查学习方法是第一个一体化图像恢复算法，不需要先验知识，不需要一个定制的网络架构，这给网络能够解决多个退化类型通过迭代学习策略

本文提出了一种轻量级高效的图像恢复通用模型，称为SimpleIR。SimpleIR可以对UHD图像进行全分辨率推理。4K，8K)在一个单一的消费级GPU上。

广泛的比较验证了SimpleIR的有效性，通过在各种图像恢复任务上实现最先进的性能，包括低光照图像增强、图像去雨、去模糊和去雪。

Review Learning

本文提出了“Review Learning”，旨在找到一套统一的预训练权重来存档通用图像恢复。在本文的训练范式中，来自前一个训练阶段的具有挑战性的样本将在后续阶段进行回顾。此外，具有挑战性的样本的数量明显小于主要的训练数据集。这确保了在避免信息污染的同时，也可以有效地防止灾难性的遗忘。在本节中，本文将详细介绍“Review Learning”和轻量级的通用UHD图像恢复模型。

理论基础

在本文的“Review Learning”方法论的核心在于持续学习的概念。这一理论表明，一个有效的学习系统应该逐步积累知识，类似于人类的学习，在人类的学习中，新的信息被吸收，而不侵蚀对先前经验的回忆（Parisi et al. 2019）。在神经网络范式中，这往往因灾难性遗忘现象而变得复杂（法语1999），在这种现象中，网络容纳新数据的能力可能以牺牲其存储的记忆为代价。到目前为止，本文设计了一个Review Learning范式，结合了一步一步学习的想法和避免灾难性遗忘的策略。Review Learning的设计如算法1所示，Review Learning的整个实现过程包括三个步骤，介绍如下：

1.基于退化复杂度的训练管道：根据退化数据集的内在学习难度对其训练顺序进行排序，这是由原始图像和退化图像之间的熵差决定的。这个排名促进了一种结构化的学习方法，从最简单的任务开始到最复杂的任务。

2.挑战性样本的获取：在初始阶段，挑战性样本通过其超过局部平均值的损失值来识别，这表明了模型的难度。从第二阶段开始，焦点转移到熵差最高的样本上，避免了混合数据集的分布，这可能会歪曲真正具有挑战性的样本。

3.顺序训练和回顾集成：模型在数据集上进行顺序训练，从第一个模型重新开始。在每个阶段之后，具有挑战性的样本将被识别和存档。随着训练对后续数据集的进展，这些具有挑战性的样本的选择被重新整合，以确保学习知识的持续强化和减轻灾难性遗忘。

在本节中，本文将详细介绍四种类型的退化的UHD图像的恢复过程，以帮助解释“Review Learning”的有效性。

如图3所示，本文展示了四种类型的UHD退化数据集之间的熵差的直方图。本文根据难度水平策略性地规划模型的学习路径，从减少到去模糊、脱色，最后到弱光增强。首先，在UHD-Snow数据集上训练模型，并根据超过周围时期平均损失的异常高损失值收集具有挑战性的样本。在预训练（第1阶段）之后，本文选择了所有具有挑战性的样本，占UHD-Snow数据集的10%。

在图4中，可以观察到在阶段1后的训练损失波动剧烈。因此，本文选择了具有前20%熵差的具有挑战性的样本，而不考虑它们的损失函数性能。此外，为了避免“信息污染”，本文将前一阶段的困难样品数量减少50%，每个后续阶段。

SimpleIR架构

Hybrid Attention Blocks：为了捕捉图像的低分辨率特征，本文提出了混合注意块（HAB）来赋予低分辨率空间更具有代表性的特征。每个HAB包含一个双流注意（DSA）和一个局部细节注意模块（LDAM）。

DualStream Attention：该模型中主要的计算开销来自于自注意层。最近已经提出了许多方法来降低自注意的二次复杂度，这是基于变压器的架构中的计算瓶颈。为了寻找适合实时移动视觉应用程序的高效和可扩展的注意机制，本文引入了一种新的方法，即选择了一种计算效率更高的策略，该策略保留了在输入数据中捕获全局依赖关系的本质。本文的双元流注意的核心是生成查询和关键向量的双重路径，设计用于封装输入特征图的交叉通道和空间上下文。

Local Detail Attention Module：受（Szegedy et al. 2016）设计的启发，该设计能够有效地捕获不同尺度上的信息，本文引入了一个轻量级的局部细节注意模块（LDAModule），以增强高分辨率图像中细粒度特征的提取。

Feed-Forward Network：为了将特征转换为一个紧凑的表示形式，本文引入了FFN模块，它由一个3×3卷积、一个1×1卷积和一个GELU函数组成。

Feature Iteration Block

实验结果