IPT - yubo105139/paper GitHub Wiki

[TOC]

论文名:Pre-Trained Image Processing Transformer

北大华为的提出的预训练模型IPT:利用ImageNet训练了一个能降噪、超分和去雨的图像预训练模型（IPT）。在去雨任务上以1.6dB超越已有最佳方案.

摘要

在这篇论文中，作者对low-level计算机视觉任务（比如降噪、超分、去雨）进行了研究并提出了一种新的预训练模型：IPT(image processing transformer)。为最大挖掘transformer的能力，作者采用知名的ImageNet制作了大量的退化图像数据对，然后采用这些训练数据对对所提IPT模型进行训练。

IPT-----有多头、多尾以适配多种退化降质模型

此外，作者还引入了对比学习以更好的适配不同的图像处理任务。微调后，预训练模型可以有效的应用不到的任务中。仅仅需要一个预训练模型，IPT即可在多个low-level基准上取得优于SOTA方案的性能。

性能图

上图给出所提方案IPT与HAN、RDN、RCDNet在超分、降噪、去雨任务上的性能对比，IPT均取得0.4-2.0dB不等的性能提升。

IPT 架构

上图给出了IPT的整体架构示意图，可以看到它包含四个成分：

用于从输入退化图像提取特征的Heads；
encoder与decoder模块用于重建输入数据中的丢失信息；
用于输出图像重建的Tails。

首先需要把feature map分块，每个patch则视为一个word。但是不同的是，因为IPT是同时训练多个task，因此模型定义了多个head和tail分别对应不同的task。

Heads

为适配不同的图像处理任务，作者提出采用多头(multi-head每个头包含三个卷积层)架构以分别处理不同的任务。假设输入图像表示为$x∈R^{3×H×W}$，每个头可以生成特征$f_H∈R^{C×H×W}$(注：C默认64)。故多头架构的计算过程可以描述为：

$f_H=H^i(x),i∈(1,...,N_t)$

其中$N_t$表示图像处理任务个数。

Transformer encoder

在将前述特征送入到Transformer之前，需要将输入特征拆分成块，每个块称之为“word”。具体来讲，输入特征$f_H∈R^{C×H×W}$将被拆分为一系列块$f_{p_i}∈R^{p^2×C},i={1,...,N}$，其中$N=HW/{P^2}$表示特征块的数量，P表示特征块尺寸。

为编码每个块的位置信息，作者还在encoder里面添加了可学习的位置编码信息$E_{p_i}∈R^{P^2×C}$。这里的encoder延续了原始Transformer，采用了多头自注意力模块和前向网络。

encoder的输出表示为$f_{E_i}∈R^{P^2×C}$,它与输入块尺寸相同，encoder的计算过程描述如下：

其中LN是layer normalization，MSA(self-attention)表示多头自注意力模块，l表示encoder的层数，FFN(feed forward network)表示前馈前向网络(它仅包含两个全连接层)。

Transformer decoder

decoder采用了与encoder类似的架构并以encoder的输出作为输入，这些task embedding是可训练的，不同的task embedding代表处理不同的task。它包含两个MSA与1个FFN。它与原始Transformer的不同之处在于：采用任务相关的embedding作为额外的输入，这些任务相关的embedding用于对不同任务进行特征编码。decoder的计算过程描述如下：

其中$f_{D_i}∈R^{P^2×C}$表示decoder的输出。decoder输出的N个尺寸为$P^2×C$的块特征将组成特征$f_D∈R^{C×H×W}$。

$f_{E_i}$是指encoder的输出，$f_{D_i}$是指decoder的输出。

Tails

Tails与Heads是相对应的，但是不同的tail的输出的大小可能不一样，例如超分，做超分时输出比输入的图像大。计算过程可以描述如下:

$f_T=T^i(f_D),i＝{1,...N_t}$

最终的输出$f_T$即为重建图像,其尺寸为$3×H^{\prime}×W^{\prime}$.输出图像的额尺寸受任务决定,比如x2超分而言,$H^{\prime}=2H,W^{\prime}=2W$.

Loss

IPT训练过程中的监督损失采用了常规的L1损失

$\mathcal{L}_{\text {supervised }}=\sum_{i=1}^{N_{t}} L_{1}\left(I P T\left(I_{\text {corrupted }}^{i}\right), I_{\text {clean }}\right)$

除了上述监督学习方式外,作者还引入对比学习以学习更通用特征以使预训练IPT可以应用到未知任务。对于给定输入$X_j$(随机从每个batch中挑选)，其decoder输出块特征描述为$f{_{D_i}^j}∈R^{P^2×C},i={1,...,N}$。作者期望通过对比学习最小化同一图像内的块特征距离，最大化不同图像的块特征距离，这里采用的对比学习损失函数定义如下：

$l\left(f_{D_{i_{1}}}^{j}, f_{D_{i_{2}}}^{j}\right)=-\log _{\sum_{i_{1}}}^{B} \exp \left(d\left(f_{D_{i_{1}}^{j}}, f_{D_{i_{2}}}^{j}\right)\right)$

$\mathcal{L}_{\text {constrastive }}=\sum_{B N^{2}} \sum_{i_{1}=1}^{N} \sum_{i_{2}=1}^{N} \sum_{j=1}^{B} l\left(f_{D_{i_{1}}}^{j}, f_{D_{i_{2}}}^{j}\right)$

其中$d(a, b)=\begin{gathered}a^{T} b \\ \|a\|\|b\|\end{gathered}$表示cosine相似性。为更充分的利用监督与自监督信息，作者定义了如下整体损失：

$\mathcal{L}_{I P T}=\lambda \cdot \mathcal{L}_{\text {constrastive }}+\mathcal{L}_{\text {supervised }}$

Lsupervised是指IPT的输出与label的L1 loss。