UNIVERSAL IMAGE RESTORATION PRE-TRAINING VIA DEGRADATION CLASSIFICATION

摘要

本文提出了退化分类预训练（DCPT）方法，该方法使模型能够学习如何对输入图像的退化类型进行分类，以用于通用图像恢复预训练。与现有的自监督预训练方法不同，DCPT将输入图像的退化类型用作一种极其微弱的监督，这种监督很容易获得，甚至在所有图像恢复数据集中都是内在存在的。DCPT包括两个主要阶段。首先，从编码器中提取图像特征。随后，利用一个轻量级解码器（如ResNet18），仅根据第一阶段提取的特征对输入图像的退化类型进行分类，而不使用输入图像。编码器通过一种简单却强大的DCPT进行预训练，该方法用于处理通用图像恢复并取得出色的性能。在采用DCPT之后，卷积神经网络（CNNs）和Transformer都表现出性能提升，在10D一体化恢复任务中增益高达2.55 dB，在混合退化场景中增益高达6.53 dB。此外，之前的自监督预训练方法（如掩码图像建模）在预训练后会丢弃解码器，而我们的DCPT能更有效地利用预训练参数。这种优势源于DCPT过程中获得的退化分类器，它有助于在针对不同退化类型训练的相同架构模型之间进行迁移学习。源代码和模型可在https://github.com/MILab-PKU/dcpt获取。

引言

尽管这些方法通过使用精确有效的提示实现了高性能，但它们未能利用恢复模型自身所蕴含的潜在先验信息。
自监督预训练策略激活神经网络内的潜在判别信息，从而有助于获取通用输入信号表示，并使预训练模型适用于下游任务。
对比学习（Chen等人，2020b；He等人，2020）通过使用对比损失（Oord等人，2018）最大化同一样本的多个增强视图之间的一致性来发现表征，从而获得具有细粒度判别信息的特征。掩蔽图像建模将BERT的成功从语言领域扩展到视觉Transformer和卷积神经网络（CNNs）。
MIM通过极高的掩码率（60%~75%）提出了一项具有挑战性的图像重建任务，强迫模型挖掘图像的内在分布。然而，用于图像修复的自监督预训练较为匮乏，仅限于单任务应用，并且没有利用大规模预训练得到的通用表征。在修复模型中找到判别性信息并采用预训练策略来增强它是至关重要的。这将创建一个预训练的恢复模型，能够处理通用的恢复任务，包括多任务（一体化）、单任务以及混合退化场景。
在本文中，我们提出退化分类能力是恢复模型中一种常被忽视但强大的固有判别信息，并且我们将其应用于通用图像恢复任务的预训练中。
我们首先探究已有的经典和一体化图像恢复架构所固有的退化分类能力。令人惊讶的是，随机初始化的模型展现出了初步的退化分类能力，在进行一体化恢复训练后，这种能力进一步增强，使它们能够识别以前从未见过的退化类型。网络本身具备退化识别能力，并且这种能力在训练阶段会逐步得到优化。受这一观察的启发，我们提出了用于通用图像恢复任务的退化分类预训练（DCPT）框架。这种方法通过预先进行退化分类学习阶段，为模型赋予了关于退化理解的强大先验信息，从而增强了模型的通用恢复能力。
具体来说，DCPT遵循编码器 - 解码器设计，其中编码器将输入图像转换为抽象特征，而一个轻量级解码器根据编码器的输出特征对退化类型进行分类。预训练的编码器在针对下游任务的微调过程中初始化一个恢复模型，从而带来显著的性能提升。
实验结果表明，DCPT框架在恢复任务领域显著优化了各种架构的性能，涵盖了一体化和混合退化条件。此外，预训练的解码器有助于在针对不同退化情况训练的相同架构模型之间进行迁移学习，从而增强模型对不同类型退化的泛化能力。

动机

期望模型能够学习针对各种退化的有效解决方案，并自主决定如何修复输入图像。我们认为，图像修复模型本质上具备区分各种退化的能力，且这种能力可以通过修复训练进一步增强。
我们进行了初步实验以作验证。提取恢复头之前的输出特征，并采用 k 近邻（kNN）分类器，基于此特征对输入图像的五种退化类型进行分类，这五种退化类型包括雾霾、雨、高斯噪声、运动模糊和低光照。我们对经典图像修复网络（Liang 等人，2021 年；Zamir 等人，2022 年；Chen 等人，2022 年）和专门的一体化任务网络（Potlapalli 等人，2023 年）进行了实验，评估了随机初始化的模型以及在三种不同（3D）退化（雾霾、雨、高斯噪声）一体化图像修复任务上训练的模型。需要提及的是，用于分类的五个目标类别涵盖了训练阶段所使用的三种退化情况。
随机初始化的模型在退化分类上的准确率可达 52%~71%。经过 3D 一体化训练后，模型在退化分类（包括未见过的退化类型）上的准确率达到 94% 及以上。图 1 绘制了 PromptIR（Potlapalli 等人，2023 年）的特征在随机初始化和 3D 一体化训练后对五种退化情况的 T-SNE 结果，从输入图像的退化角度更直观地展示了恢复模型的分类能力。

这些结果得出了三个有趣的结论：(1) 随机初始化的模型表现出对退化进行分类的内在能力。(2) 在一体化任务上训练的模型展现出辨别未知退化的能力。(3) 在恢复模型的早期训练中有一个退化理解步骤。
这表明图像修复模型本身具有对各种退化类型进行分类的能力。在一体化修复模型的训练过程中，可以推测，在模型执行修复任务的同时，它也在同步接受训练以识别输入图像中存在的退化类型。
为确保卓越的恢复性能，恢复模型在训练前具备足够的退化分类能力至关重要。

DCPT

基于这个简单明了的想法，我们提出了退化分类预训练（DCPT）。
DCPT 由一个编码器和解码器组成，编码器包含去除了恢复模块的恢复模型（Liang 等人，2021 年；Zamir 等人，2022 年；Chen 等人，2022 年；Potlapalli 等人，2023 年），解码器则根据编码器提取的特征对输入图像的退化情况进行分类。在退化分类（DC）阶段，给定具有退化$D_{gt}$的输入图像$x_{degrad}$，编码器的特征$F$被输入到解码器中。由于解码器的作用是分类而非图像重建，其设计借鉴了经典的分类网络（He等人，2016年），而非经典自动编码器式预训练方法中常见的架构。在生成阶段，编码器的特征F被输入到恢复模块中，以保留编码器的生成能力。图2展示了这一整体流程。

Extract multi-level features

为了实现更有效的退化分类，从包含更丰富高级语义信息的更深层次中提取特征至关重要。仅依赖最深层的特征来计算损失函数，可能会由于特征提取过程中编码器的长残差连接丢失（Zhang等人，2017），导致较浅层出现梯度消失问题。为了实现平衡，从编码器后半部分的每个模块中提取特征。

退化分类

经过多层次特征提取后，我们将{Fi}输入到一个轻量级解码器中，以对输入图像的退化情况进行分类。为了更好地聚合提取的特征，有必要对特征 {Fi} 进行缩放。缩放系数{ωi}是可学习的。然后，将缩放后的特征F′i = ωiFi输入到ResNet18的第i个倒数层模块中，以对退化情况进行分类。 为了稳定训练过程（Liu等人，2022），我们将解码器中的归一化层从批归一化（Ioffe和Szegedy，2015）替换为层归一化（Ba等人，2016）。

获取图像恢复数据的困难（Li 等人，2023b）导致了代表不同类型退化的数据集数量不均衡。例如，去雨数据集Rain200L（Yang等人，2017年）仅包含200张图像，而图像去雾数据集RESIDE（Li等人，2018年）包含72135张图像。这种不平衡在对退化进行分类时带来了重大的长尾挑战。为了解决这个问题，我们采用焦点损失（Focal Loss，Lin等人，2017年）作为长尾退化分类的损失函数。

保留生成能力

必须记住，编码器的目标是将低质量图像恢复为高质量图像。确保预训练的编码器保持其生成能力很重要。与经典生成方法（Kingma和Welling，2013年）的训练过程类似，在编码器之后添加一个卷积层，使其能够从特征Fl中重建出干净的xˆ，如图2（b）所示。

将一次预训练迭代分为两个交替阶段。

在DCPT过程中，损失函数Lpix和Lcls都需要利用编码器生成的特征。同时执行退化分类阶段和生成阶段会导致编码器同时接收两个不同的梯度流，这是不可行的。为了解决这个问题，我们在一次预训练迭代中交替进行这两个阶段。

DC（退化分类）引导训练

与以往的预训练方法（He等人，2022年；Liu等人，2023a）不同，预训练后解码器不会被丢弃。这带来了用于跨退化迁移学习的DC引导训练。如图3所示，在恢复干净图像${\hat{x}}_{restore}$时，有必要确保恢复模型的特征能够识别输入图像的退化情况。在DC引导的训练中，解码器的冻结参数由直流预训练（DCPT）设置。
在DC引导训练中，整体损失函数为：

交叉熵损失被用作DC引导训练的分类损失，因为它在迁移学习条件下具有适用性，特别是针对单一类型的退化。解码器的作用仅限于将输入图像分为两类：清晰的和退化的。配对数据为两个类别提供了数量相等的样本，不存在类别不平衡问题，因此可以使用交叉熵损失，而不是焦点损失。

实验

Dataset Sampler in all-in-one setting

对于训练数据较少的退化任务，如图像去雨，我们使用重复采样器技术，以确保每种退化都有足够的训练对。重复比例为[1H, 120R, 9N]，其中H、R和N分别代表去雾、去雨和去噪。对于5D一体化图像恢复，将去雾、去雨、高斯去噪、运动去模糊和低光照增强的数据集整合起来进行微调。重复比例为[1H, 300R, 15N, 5B, 60L]，其中H、R、N、B和L分别代表去雾、去雨、去噪、去模糊和低光照增强。

DCPT设置

在DCPT期间，图像恢复模型（编码器）和退化分类器（解码器）均由AdamW（Kingma和Ba，2014年）进行训练，在4块NVIDIA L40 GPU上，对128×128的图像块，以批量大小32，无权重衰减地训练10万次迭代。由于采用了异构的编码器 - 解码器设计，我们对编码器和解码器采用不同的学习率。编码器的学习率设置为3×10⁻⁴，解码器的学习率设置为1×10⁻⁴。在DCPT期间，学习率保持不变。DCPT之后，编码器的参数将用于初始化图像恢复模型。

一体化设置细节

我们使用DCPT预训练的参数初始化图像恢复模型。为了公平和方便起见，我们对不同的主干网络采用相同的训练策略。我们使用AdamW（Kingma和Ba，2014年）优化器，初始学习率为3×10⁻⁴ ，并通过余弦退火调度逐渐降低到1×10⁻⁶ 来训练我们的图像恢复模型。训练在4块NVIDIA L40 GPU上进行，批量大小为32，共运行75万次迭代。