DiffIR: Efficient Diffusion Model for Image Restoration

• 清华&ETH&字节&UTD
• https://github.com/Zj-BinXia/DiffIR

问题引入

• IR任务和image synthesis任务不同点是IR任务本身有一个很强的低质量图片作为先验，所以可以不完全遵循图片生成的范式，本文主要在compact的IPR空间进行DM；
• 本文提出的模型分为三个部分，1)CPEN(compact IR prior extraction network)来得到IPR(IR prior representation)，这个作为回归模型的指导信息；2)DIRformer回归模型，类比为decoder；3)DM来通过LQ图片得到IPR
• 训练分为两个stage，首先第一个stage训练CPEN和DIRformer，此时CPEN输入的是高质量图片；第二个stage使用的IPR是DM得到的；

methods

• stage1: 训练CPEN和DIRformer，首先将gt和LQ concat到一起，然后经过pixelunshuffle得到CPEN的输入，输出IPR $Z=CPE{N}_{S1}(PixelUnshuffle(Concat(I_{GT},I_{LQ}))),Z\in {\mathbb{R}}^{4C^{\prime }}$，之后IPR被送到DIRformer的DGFN和DMTA模块，第一阶段训练的损失是GT和生成HQ的L1损失，超分和inpainting任务还有erceptual loss and adversarial
  loss；
• DMTA的操作$F^{\prime }=W_l^{1}Z\odot Norm(F)+W_l^{2}Z$，其中$W_l$是linear层，$F,F^{\prime }$分别是输入和输出的feature map，$Q=W_d^Q{W}_c^Q{F}^{\prime },K=W_d^K{W}_c^K{F}^{\prime },V=W_d^V{W}_c^V{F}^{\prime }$，其中$W_d$是depthwise卷积，$W_c$是pointwise卷积，之后被reshape成$\hat{Q}\in {\mathbb{R}}^{\hat{H}\hat{W}\times \hat{C}},\hat{K}\in {\mathbb{R}}^{\hat{C}\times \hat{H}\hat{W}},\hat{V}\in {\mathbb{R}}^{\hat{H}\hat{W}\times \hat{C}}$，最后$\hat{F}=W_c\hat{V}\cdot Softmax(\hat{K}\cdot \hat{Q}/\gamma )+F$；
• DGFN的操作：$\hat{F}=GELU(W_d^1{W}_c^1{F}^{\prime })\odot W_d^2{W}_c^2{F}^{\prime }+F$；
• stage2：同时训练三个部分，首先使用$CPE{N}_{S1}$得到$Z$，之后经过diffusion process得到$Z_T\in {\mathbb{R}}^{4C^{\prime }}$，$CPE{N}_{S2}$得到$D=CPE{N}_{S2}(PixelUnshuffle(I_{LQ}))$，之后进行DM，以D为条件，进行去噪t-1次得到$\hat{Z}$，和$CPE{N}_{S1}$得到的$Z$计算损失$L_{diff}=\frac{1}{4C^{\prime }}{\sum }_{i=1}^{4C^{\prime }}|\hat{Z}(i)-Z(i)|$，这损失和stage1的损失在一起计算总损失；

实验