超分算法IPT：Pre-Trained Image Processing Transformer

本文是一个基于transformer的预训练通用模型，针对低级视觉任务还没有人提出预训练的模型，所以作者使用了超大数据集训练出了image processing transformer (IPT)。可以微调后应用于图像重建、去噪、去雨等等。具体结构作者使用了一个多头多尾共享躯干的结构。应对不同的任务，有针对性不同的头部和尾部，分别使用不同的处理方式。中间是一个transformer编解码器结构。将头部输出的特征图像unfold成”词向量“形式和位置嵌入相加后输入encoder，encoder是常规的结构，包括一个LN和MSA接残差、LN和FFN接残差（FNN有两层全连接）。解码器结构和常规transformer也差不多，但是多了一个特定任务标签嵌入，加入在decoder输入第一次MSA的Q和K上、第二次Q的上。尾部也是多个针对特定任务的结构，用于还原图像尺寸。

数据集是使用ImageNet自己处理的，因为需要大量数据才能训练出较好的预训练模型。

LOSS使用有监督的L1L损失和对比学习损失函数。

原文链接：IPT：Pre-Trained Image Processing Transformer
源码地址：
https://github.com/huawei-noah/Pretrained-IPT
and
https://gitee.com/mindspore/mindspore/tree/master/model_zoo/research/cv/IPT

Abstract

随着现代硬件计算能力的强劲增长，在大规模数据集上学习的预训练的深度学习模型（如BERT、GPT-3）已显示出其比传统方法更有效。这一巨大进步主要归功于transformer及其变体架构的表示能力。

在本文中，作者试图在低级计算机视觉任务（例如去噪、超分辨率和去雨），也开发一个预训练模型，即图像处理Transformer——image processing transformer （IPT）。为了最大限度地挖掘transformer的性能，作者使用了著名的ImageNet基准生成大量损坏的图像对。IPT模型在这些图像上进行多头多尾训练。此外，为了更好地适应不同的图像处理任务，还引入了对比学习。因此，经过微调后，预训练的模型可以有效地用于所需的任务。虽然只有一个预训练的模型，IPT在各种低水平基准上的表现优于当前最先进的方法。

1 Introduction

图像处理是计算机视觉系统的低级部分的一个组成部分。图像处理的结果会在很大程度上影响后续高级部分对图像数据进行识别和理解。由于许多图像处理任务都是相关的，因此很自然地，人们会期望在一个数据集上预训练的模型可以对另一个数据集有所帮助。但很少有研究将预训练推广到图像处理任务中。

现在，在自然语言处理和计算机视觉方面进行预训练很常见。

1. 目标检测模型的主干通常是在ImageNet分类上预训练的，包括AlexNet、VGGNet和ResNet。
2. 开创性的Transformer也已广泛应用于许多自然语言处理（NLP）任务，如翻译和问答。基本都是大型文本语料库上预训练基于transformer的模型，并在特定于任务的数据集上对其进行微调。
3. Transformer的变体，如Bert和GPT-3，进一步丰富了训练数据，提高了预训练的能力。

已经有学者尝试将Transformer推广到计算机视觉领域。例如，Wang等人和Fu等人应用了基于自注意的模型来捕捉图像上的全局信息；Carion等人提出DERT使用transformer架构进行端到端目标检测；Dosovitskiy等人引入了视觉Transformer（ViT），将输入图像处理为16×16的token，并在图像识别方面取得了优异的效果。

图像处理任务的预训练方法需要解决两个问题：

1. 任务特定的数据受到限制。这一问题在涉及付费数据或数据隐私的图像处理任务中更加严重，如医学图像和卫星图像。各种不一致因素（例如摄像机参数、照明和天气）可能会进一步干扰用于训练的数据分布。
2. 在应用测试前，并不知道将用于哪种类型的图像处理任务。因此，必须准备一系列的图像处理模块。各个模块有明确的任务目标，但一些潜在的部分可以共享数据。

在本文中，作者使用transformer架构开发了一个预训练的图像处理模型，即图像处理transformer（IPT）。由于预训练模型需要与不同的图像处理任务兼容，包括超分辨率、去噪和去雨，整个网络由多对对应于不同任务的头部和尾部以及单个共享体组成。

由于需要使用大规模数据集挖掘Transformer的潜力，用大量、多样性的图像来训练IPT模型。为此，选择了ImageNet基准，它包含1000个类别的各种高分辨率图像。对于ImageNet中的每个图像，使用几种精心设计的操作来为不同的任务生成多个损坏的副本。例如，超分辨率任务的训练样本是通过对原始图像进行下采样生成的。用于IPT训练的整个数据集包含约1000多万张图像 。然后在庞大的数据集上训练transformer架构。

训练图像被输入到特定的头部，生成的特征被裁剪成小块（即“token”），然后展平成序列。Transformer处理展开的特征，编码器和解码器分别使用位置和任务嵌入。此外，根据特定任务，尾部被强制预测具有不同输出大小的原始图像。此外，为了更好地适应不同的图像处理任务，还引入了不同输入块之间关系的对比损失。提出的图像处理Transformer是以端到端的方式学习的。在多个基准上进行的实验结果表明，经过预训练的IPT模型可以在微调后显著增强，从而超过大多数现有方法。

2 Method

2.1 IPT architecture

IPT的总体架构由四个部分组成：头部用于从输入损坏图像中提取特征，编码器、解码器（用于恢复输入数据中缺失信息）以及用于将特征映射为恢复图像的尾部。

Heads：
为了应对不同的图像处理任务，使用了一种多头部结构来分别处理每个任务，其中每个头部由三个卷积层组成。
输入图像表示为$x\in R^{3\times H\times W}$，头部生成特征图$f_H\in R^{C\times H\times W}$，（通常使用C=64）。计算公式为
$$f_H=H^i（x）$$其中$H^i(i=1，\dots ，N_t)$表示第i个任务的头部，$N_t$表示任务数。

Transformer encoder:

1. 先将特征分割为patch，每个patch展开为一列向量，视为一个“词向量”。输入特征$f_H\in R^{C\times H\times W}$ reshape 成一序列patch $f_{pi}\in R^{P^2\times C}，i=1，\dots ，N$，patch数量为$N=\frac{HW}{p^2}$（序列长度），p是patch大小。
2. 再加入位置信息，每个patch都添加了可学习的位置编码$E_{pi}\in R^{P^2\times C}$。$E_{pi}+f_{pi}$将直接输入Transformer编码器。
3. Transformer encoder继承最原始结构，有一个多头自注意模块和一个前馈网络。经过encoder注意力计算不改变输入输出尺寸。encoder内部公式如下：

$l$表示编码器中的层数，MSA表示传统Transformer模型中的多头自注意模块，LN表示层规范化，FFN表示前馈网络，其中包含两个全连接层。

Transformer decoder：
Decoder也和原始的decoder无甚差别只是多加了一个任务类型嵌入。Transformer edcoder由两个多头自注意（MSA）层和一个前馈网络（FFN）组成。使用特定于任务的嵌入作为decoder的额外输入。这些特定于任务的嵌入$E_t^i\in R^{P^2\times C}，i=1，\dots ，N_t$为学习不同任务解码特征。最后将解码得到的大小为$p^2\times C$的N个特征重塑为大小为$C\times H\times W$的特征$f_D$。解码器的计算公式如下：

Tails：
尾部的性质与头部相同，使用多个尾部来处理不同的任务。计算公式为$f_T=T^i（f_D）$，其中$T_i（i=1，\dots ，N_t）$表示第$i$个任务的头部，$N_t$表示任务数。输出$f_T$是由特定任务确定的图像大小为$3\times H\prime \times W\prime$的结果。例如，对于2×超分辨率任务，H′=2H，W′=2W。

2.2 Pre-training on ImageNet

决定一个模型成功与否，除了本身的网络结构，还有一个关键因素是大规模数据集的良好使用。
相比于图像分类数据集，图像处理任务的可用数据集数量较少（例如，用于图像超分辨率任务的DIV2K数据集上只有2000张图像），作者使用著名的ImageNet作为基线数据集，自己生成所需的数据集，对IPT模型进行预训练。

由于ImageNet基准中的图像具有高度多样性，其中包含来自1000个不同类别的100多万张自然图像。这些图像具有丰富的纹理和颜色信息。首先移除语义标签，然后手动从这些未标记的图像中合成各种损坏的图像，并针对不同的任务使用各种降级模型。

1. 超分辨率任务通常采用双三次下采样生成低分辨率图像，
2. 去噪任务在不同噪声水平的原始图像中加入高斯噪声生成带噪图像。

这些合成图像可以显著提高学习深度网络的性能，包括CNN和transformer结构，这将在实验部分证明。

在监督模式下IPT的损失函数可以表示为：

L1分别表示重建图像的常规L1损失，$I_{corrupted}^i$任务$i$的损坏图像。此外，公式(4)意味着所提出的框架同时接受多个图像处理任务的训练。

训练过程：
具体来说，对于每个批次，从{N_t}个任务中随机选择一个任务进行训练，每个任务将同时使用相应的头、尾和任务嵌入进行处理。在对IPT模型进行预训练后，它将捕获大量图像处理任务的内在特征和转换。因此只需进一步微调，使用对应任务提供的数据集就可以应用于特定的任务。此外，为了节省计算成本，在微调时将删除其他头部和尾部，并根据反向传播更新剩余的头部、尾部和主干中的参数。

额外的损失函数：
然而，由于退化模型的多样性，无法为所有图像处理任务合成图像。并且在实践中可能存在各种各样的噪声。为此，应进一步增强生成IPT的泛化能力。

与预训练的自然语言处理模型类似，图像块之间的关系也提供了信息。图像场景中的patch可以看作是自然语言处理中的一个词。因此，引入对比学习来学习通用特征，使预训练的IPT模型可以用于未知任务。对比学习的目标是最小化来自相同图像的patch特征之间的距离，同时最大化来自不同图像的patch之间的距离。对比学习的损失函数公式如下：

$d(a,b)$表示余弦相似度。此外，为了充分利用监督和自监督信息，$L_{IPT}$作为IPT的最终目标函数。λ用来平衡对比损失与监督损失。损失函数公式重新整合为：

3 Experiments

数据集：
ImageNet数据集，该数据集由超过100万张具有高度多样性的彩色图像组成。训练图像被裁剪成48×48块，有3个通道用于训练。有超过10M的patch用于训练IPT模型。生成了6种退化类型的受损图像：2×、3×、4×双三次插值下采样、30、50噪声级高斯噪声和添加雨纹。使用32张Nvidia Nvidia Tesla V100卡。

3.1 Super-resolution

将IPT模型与几种最先进的基于CNN的SR方法进行了比较。如表1所示，预训练IPT优于所有其他方法，并在所有数据集的×2、×3、×4尺度上实现了最佳性能。值得强调的是，IPT在×2尺度的Urban100数据集上实现了33.76dB的峰值信噪比，这比其他方法的峰值信噪比高出∼0.4dB，而之前的SOTA方法与其他方法相比只能实现<0.2dB的改进，这表明该模型利用大规模预训练的优越性。

图3在Urban100数据集上以4×比例展示了模型的可视化结果。高比例因子下恢复原始高分辨率图像很困难，因为会导致大量信息丢失。以前的方法生成模糊图像，而IPT模型生成的超分辨率图像可以很好地从低分辨率图像中恢复细节。

3.2 Denoising

IPT与各种最先进的模型进行比较。表2展示了BSD68和Urban100数据集上的彩色图像去噪结果。
在不同高斯噪声水平下，IPT在所有去噪方法中取得了最好的效果。此外，IPT模型在Urban100数据集上超过SOTA方法~2dB，这证明了预训练的有效性和基于Transformer的模型的优越性。

图4显示了结果图像的可视化。如图所示，噪声图像很难识别，很难恢复干净的图像。现有的方法无法重建足够的细节并生成异常像素。而预训练的模型可以很好地恢复这只猫（？？？）头发中的一些细节（三只小小鸟），视觉质量明显优于之前的所有模型。

3.3 Deraining

对于图像去雨任务，在合成的Rain100L数据集上评估了IPT模型，该数据集由100幅RAIN图像组成。定量结果见表3。与最先进的方法相比，IPT实现了最好的性能（41.62dB），提高了1.62dB。图5显示了可视化结果。以往的方法由于缺乏图像先验知识，无法重建原始的干净图像。但IPT模型可以呈现与真实图像完全相同的视觉效果。在视觉质量方面通过了之前的所有算法。这一结果证实了所提出模型的普遍性。

3.4 Generalization Ability

虽然作者生成了各种损坏的图像，但自然图像的复杂性很高，无法合成所有可能的图像来预训练Transformer模型。然而，一个好的预训练模型应该能够很好地适应NLP领域的其他任务。为此，进行了一些实验来验证模型的泛化能力。在实验中，测试了合成ImageNet数据集中未包含的损坏图像，即分别使用10级和70级噪声进行图像去噪。使用对应的头部和尾部作为预训练模型进行图像去噪任务。详细结果如表4所示，比较了使用预训练的IPT模型和最先进的图像去噪方法的性能。显然，IPT模型优于其他常规方法，表明预训练的模型可以从大规模数据集中捕获更多有用的信息和特征。

3.5 Ablation Study

①数据百分比的影响：
使用合成的ImageNet数据集的20%、40%、60%、80%和100%的百分比来分析所用数据数量对结果性能的影响。图6显示了不同预训练模型的结果。当模型没有经过预训练或使用少量（<60%）数据集进行预训练时，CNN模型可以获得更好的性能。相比之下，当使用大规模数据时，基于Transformer的模型压倒了CNN模型，这证明了IPT预训练模型的有效性。

②对比学习的影响：

为了提高预训练模型的表征能力，将对比学习损失嵌入到训练过程中。用Set4数据集评估其在×2尺度超分辨率任务中的有效性。表5显示了超参数λ对平衡两项损失函数的影响。当λ=0时，IPT模型仅使用监督学习方法进行训练，得到的PSNR值为38.27dB。当采用对比损失进行自监督学习时，该模型可以获得38.37dB的PSNR值（λ=0.1），比用λ=0训练的模型高出约0.1dB。这些结果进一步证明了对比学习对预训练IPT模型的有效性。

4 Conclusion

本文旨在利用预训练的Transformer模型（IPT）解决图像处理问题。IPT模型设计有多个头部、多个尾部和一个共享的Transformer体，用于服务不同的图像处理任务，如图像超分辨率和去噪、去雨。

为了最大限度地挖掘transformer架构在各种任务上的性能，探索了一个综合的ImageNet数据集。其中，每个原始图像将被降级为一系列对应的成对训练数据。

然后使用有监督和自监督的方法对IPT模型进行训练，这些方法显示了捕获底层图像处理固有特征的强大能力。

实验结果表明，经过快速微调后，IPT仅使用一个预训练的模型，就可以超越最先进的方法。在未来的工作中，还可以把IPT模型扩展到更多的任务，如图像修复、去雾等。

最后祝各位科研顺利，身体健康，万事胜意~