【论文精读】Learning Texture Transformer Network for Image Super-Resolution

出处

2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)

贡献

1.提出了一种包含四个密切相关模块的纹理转换器，用于图像超分辨率重建。
2.提出了一种新的跨尺度特征集成模块，用于图像生成任务，使我们的方法能够通过堆叠多个纹理转换器来学习更强大的特征表示。

纹理转换器

提出了用于图像超分辨率的纹理变换网络(TTSR)。在纹理转换器的基础上，提出了跨尺度特征集成模块(CSFI)来进一步提高模型的性能。

纹理转换器的结构如上图所示，$LR$、$LR\uparrow$和$Ref$分别表示输入图像、4次双三次插值法对输入图像进行上采样和参考图像。并对$Ref$进行4次双三次插值下采样和上采样，得到$Ref\downarrow \uparrow$与$LR\uparrow$结构保持一致。纹理转换器将$LR$、$Ref$、$Ref\downarrow \uparrow$和$LR\uparrow$作为输入，输出合成的特征图，进一步用于生成$HR$预测。纹理转换器包括4部分：可学习纹理提取器（LTE）、相关性嵌入模块（RE）、特征转移的硬注意模块（HA）和特征合成的软注意模块（SA）。

1.可学习纹理提取器（LTE）

该部分设计了一个可学习的纹理提取器，它的参数会在端到端的训练过程中进行更新。这样的设计对$LR$和$Ref$图像进行联合特征学习，其中可以捕获更精确的纹理特征。纹理提取过程可以表示为：
$$Q=LTE(LR\uparrow )\text{\hspace{1em}(1)}$$
$$K=LTE(Ref\downarrow \uparrow )\text{\hspace{1em}(2)}$$
$$V=LTE(Ref)\text{\hspace{1em}(3)}$$
其中$LTE(\cdot )$表示我们的可学习纹理提取器的输出。提取的纹理特征Q、K和V表示转换器内部注意力机制的三个基本元素，并将进一步用于相关性嵌入模块。

2.相关嵌入（RE）

相关性嵌入的目的使通过估计Q和K之间的相似度来嵌入$LR$和$Ref$图像之间的相关性。将Q和K都展开成块，$q_i(i\in [1,H_{LR}\times W_{LR}])$和$k_j(j\in [1,H_{R}ef\times W_{Ref}])$。然后对每一个块$q_i$和$k_j$，通过归一化内积来计算两个块之间的相关性$r_{i,j}$。
$$r_{i,j}=<\frac{q_i}{||q_i||},\frac{k_j}{||k_j||}>\text{\hspace{1em}(4)}$$
利用相关性得到硬注意图和软注意图。

3.硬注意（HA）

这里提出了一个硬注意模块来转移参考图像中的HR纹理特征V，在硬注意力模块中对于每个查询$q_i$只从V中最相关的位置转移特征。
具体来说，首先计算硬关注图H，其中第i个元素$h_i(i\in [1,H_{LR}\times W_{LR}])$是由相关性$r_{i,j}$得来。
$$h_i=\underset{j}{\mathrm{argmax}}{r}_{i,j}\text{\hspace{1em}(5)}$$
$h_i$的值为硬索引，表示参考图像和LR图像中第i个最相关位置。为了从参考图像中获取纹理特征T，使用硬注意图作为索引，对V的展开块进行索引选择。
$$t_i=v_{h_i}\text{\hspace{1em}(6)}$$
根据该输出结果，我们得到了LR图像和HR图像的特征表示T，它将进一步用于软注意模块。

4.软注意（SA）

提出的软注意模块将从DNN主干的LR图像和HR纹理特征T和LR特征F中合成特征。在合成过程中，为加强相关的纹理传递，从$r_{i,j}$中计算软注意图S，以表示T中每个位置的所传送纹理特征的置信度。
$$s_i=\underset{j}{\max }{r}_{i,j}\text{\hspace{1em}(7)}$$
其中侧重软注意图S的第i个位置，而不是直接将关注图S应用到T。首先融合HR纹理特征T和LR特征F，以利用来自LR图像的更多信息。融合后的特征被进一步以元素方式乘以软注意图S，并加回F，以获得纹理转换器的最终输出。
$$F_{out}=F+Conv(Concat(F,T))\odot S\text{\hspace{1em}(8)}$$
其中，$F_{out}$表示合成的输出特征，$Conv$和$Concat$分别表示卷积层和拼接操作，$\odot$表示元素相乘。
综上所述，纹理转换器可以有效地将参考图像中的相关HR纹理特征转换为LR特征，从而提高纹理生成的精确度。

跨尺度特征融合

纹理转换器可以通过跨尺度融合模块以跨尺度的方式进一步堆叠。该体系结构如上图所示，3个纹理转换器输出3个分辨率不同的纹理特征融合LR图像中。
提出的跨尺度的特征融合模块（CSFI）在不同尺度的特征之间交换信息，每次将LR特征向上采样到下一个比例时，都会应用CSFI模块。对于CSFI模块中的每个尺度，通过上/下采样接受来自其他尺度的交换特征，然后在通道维度中进行串联操作。接着，卷积层将特征映射到原始数量的通道中。在这样的设计中，从堆叠的纹理转换器传输的纹理特征在每个尺度上交换，这实现了更强大的特征表示。这种跨尺度的特征集成模块进一步提高了方法的性能。

损失函数

总体的损失表示为：
$${\mathcal{L}}_{overall}={\lambda }_{rec}{\mathcal{L}}_{rec}+{\lambda }_{adv}{\mathcal{L}}_{adv}+{\lambda }_{per}{\mathcal{L}}_{per}\text{\hspace{1em}(9)}$$

重建损失

$${\mathcal{L}}_{rec}=\frac{1}{CHW}||I^{HR}-I^{SR}|{|}_1\text{\hspace{1em}(10)}$$

对抗性损失

采用了WGAN-GP，它提出的梯度范数的惩罚来代替权重裁剪，使得训练更稳定，性能更好。
$${\mathcal{L}}_D=\underset{\tilde{x}\backsim {\mathbb{P}}_g}{E}[D(\tilde{x})]-\underset{x\backsim {\mathbb{P}}_r}{E}[D(x)]+\lambda \underset{\hat{x}\backsim {\mathbb{P}}_{\hat{x}}}{E}[(||{\nabla }_{\hat{x}}D(\hat{x})|{|}_2-1{)}^2]\text{\hspace{1em}(11)}$$
$${\mathcal{L}}_G=-\underset{\tilde{x}\backsim {\mathbb{P}}_g}{E}[D(\tilde{x})]\text{\hspace{1em}(12)}$$

感知损失

$${\mathcal{L}}_{per}=\frac{1}{C_i{H}_i{W}_i}||{\varphi }_i^{vgg}(I^{SR})-{\varphi }_i^{vgg}(I^{HR})|{|}_2^2+\frac{1}{C_j{H}_j{W}_j}||{\varphi }_j^{lte}(I^{SR})-T|{|}_2^2\text{\hspace{1em}(13)}$$
第一部分是传统的感知损失，其中$\varphi (\cdot {)}_i^{vgg}$表示第i层vgg19的特征图。
第二部分是传递性感知损失，其中$LTE(\cdot )$表示$\varphi$第j层提取的纹理特征图。