FADE

FADE: fusing the assets of decoder and encoder for task-agnostic upsampling

对于不同的稠密预测任务，有些任务注重于语义信息（semantic information），它们对图像位置信息敏感（region-sensitive）而有的任务（图像分割）注重于细节（detail-sensitive）。原有的上采样算子都是用于特定任务而不是通用型，本篇文章主要从encoder与decoder两方面解释了为什么CARAFE在semantic segmentation上优于IndexNet；IndexNet在image matting上优于CARAFE，并且提出了一种新的通用任务的采样算子FADE，FADE的有效性在于1.同时使用encoder与decoder，2.使用 semi-shift convolution 保证感受野大小，3.在产生上采样算子前使用一个gating mechanism模块，设计成与decoder无关以增强上采样算子的细节增强能力。

encoder&decoder

我们先提出一个假设：一个理想的上采样算子需要能够很好地保留semantic information的同时也需要补偿下采样过程中失去的细节信息；对于前者，semantic information已经嵌入在了decoder中，对于后者则存在于encoder中。我们进一步假设：正是因为encoder或者decoder中的信息使用不够充分或者有偏向性才导致采样算子有着task-dependent的特性。

下面通过采样算子的原理解释为什么现有的采样算子task-dependent：

encoder&decoder

CARAFE的上采样算子通过下采样之后的特征图产生，即由decoder feature产生，这样就充分利用了decoder feature的信息而丢失了下采样前的图像细节信息。
IndexNet和A2U的上采样核通过下采样前的feature map产生，这样就充分利用了原feature map中的细节信息，而没有利用采样后的semantic information
于是我们提出了FADE：通过将encoder与decoder的信息融合生成采样算子

那么接下来的问题是如何将两个信息全部利用起来？

How to leverage both encoder and decoder features for task-agnostic upsampling?

一个朴素的想法是将encoder与decoder所提供的信息concatenate起来，通过卷积进行特征提取（卷积核实质上就是一个滤波器）生成上采样算子，但是主要的问题是：encoder与decoder所提供的feature map大小不一样（mismatched resolution），考虑如果对decoder使用参数为2的临近插值（2NN）进行信息补全，同时上采样率为2，当我们使用 $3 \times 3$ 大小的卷积核进行卷积时，本来是 $3 \times 3$ 大小的感受野就退化为只能感受到4个像素信息，同时实验也证明这种方法并不好。于是启发我们设计一个新的解决方案：semi-shift convolution。每当卷积窗口在 encoder feature 中移动两个像素点时，decoder feature 中的卷积窗口移动一个像素点，这一点和 CARAFE 的方式是相同的

semi-shift convolution

FADE 的设计

FADE

encoder feature 与 decoder feature 先经过一个半移位卷积生成 pre-upsampled feature，如果直接使用这份 feture map 会发现它恢复细节的能力不够，这时候引入 gated mechanism 进行微调，decoder feature 通过 $1 \times 1$ 的卷积进行通道压缩，再通过一次 NN插值上采样和 sigmoid 函数生成 $\mathcal{G}$ ，这个 $\mathcal{G}$ 用于平衡 encoder feature 和 pre-upsampled 的权重，最终生成的上采样 feature map $\mathcal{F}$ 由如下公式生成：

\mathcal{F}_{upsampled} = \mathcal{F}_{encoder} \odot \mathcal{G} + \mathcal{F}_{pre-upsampled} \odot (1 - \mathcal{G})

Gate是否有必要的讨论

在后续实验中，其实可以发现 Gate 其实并不是必须的，直接取 $\mathcal{G}=1$ 并不会对算子的性能产生很大的破坏，甚至有时候 $\mathcal{G}=1$ 的效果比动态生成 $\mathcal{G}$ 的效果好，为了简单期间，后续实现的代码我就不考虑 Gate 了

H2L与L2H

两种设计方式

记 $conv_{/s}$ 为步长为 $s$ 的卷积，我们上面讨论的实现方式为对高分辨率的 encoder feature 使用步长为 2 的卷积去匹配 decoder feature 的大小，称这种实现方式为 H2L；反过来，我们也可以通过先对低分辨率 decoder feature 进行 NN 上采样，这样用低分辨率特征去匹配高分辨率特征称为 L2H

在算子效果上，H2L 方式的计算开销更大，但是效果更好；L2H 方式的计算开销更小，但是效果不如 H2L

代码实现：

为了简单期间，我只实现了 FADE 的 H2L 方式算子

import torch
from torch import nn
import torch.nn.functional as F
import einops


class SemiShift(nn.Module):
    def __init__(self, in_channels_en, in_channels_de, scale=2, embedding_dim=64, up_kernel_size=5):
        super(SemiShift, self).__init__()
        self.scale = scale
        self.embedding_dim = embedding_dim
        self.up_kernel_size = up_kernel_size
        self.conv1_en = nn.Conv2d(in_channels_en, embedding_dim, kernel_size=1)
        self.conv1_de = nn.Conv2d(in_channels_de, embedding_dim, kernel_size=1, bias=False)

        self.conv2_kernels = nn.Parameter(torch.empty((up_kernel_size ** 2, embedding_dim, 3, 3)))
        nn.init.xavier_normal_(self.conv2_kernels, gain=1)
        self.conv2_bias = nn.Parameter(torch.empty(up_kernel_size ** 2))
        nn.init.constant_(self.conv2_bias, val=0)

    def forward(self, en, de):
        B, _, H, W = de.shape
        compressed_en = self.conv1_en(en)
        compressed_de = self.conv1_de(de)
        pad_en = []
        pad_en.append(F.pad(compressed_en, pad=[1, 0, 1, 0]))
        pad_en.append(F.pad(compressed_en, pad=[0, 1, 1, 0]))
        pad_en.append(F.pad(compressed_en, pad=[1, 0, 0, 1]))
        pad_en.append(F.pad(compressed_en, pad=[0, 1, 0, 1]))
        pad_en = torch.cat(pad_en, dim=1)

        # h = H + 1, w = W + 1
        pad_en = einops.rearrange(pad_en, 'b (c scale_2) h w -> (b scale_2) c h w', scale_2=self.scale ** 2)
        kernels = F.conv2d(pad_en, self.conv2_kernels, self.conv2_bias, stride=2)
        # c = self.up_kernel_size ** 2
        kernels = einops.rearrange(kernels, '(b scale_2) c h w -> b scale_2 c h w', scale_2=self.scale ** 2)

        kernels = kernels + F.conv2d(compressed_de, self.conv2_kernels, self.conv2_bias, stride=1,
                           padding=1).unsqueeze(1)  # scale_2 维度进行广播
        
        kernels = einops.rearrange(kernels, 'b (scale1 scale2) c h w -> b c (h scale1) (w scale2)',
                                   scale1=self.scale, scale2=self.scale)
        return kernels



class FADE(nn.Module):
    def __init__(self, in_channels_en, in_channels_de=None, scale=2, up_kernel_size=5):
        super(FADE, self).__init__()
        in_channels_de = in_channels_de if in_channels_de is not None else in_channels_en
        self.scale = scale
        self.up_kernel_size = up_kernel_size
        self.kernel_generator = SemiShift(in_channels_en, in_channels_de,
                                              up_kernel_size=up_kernel_size, scale=scale)
        self.carafe = CARAFE()

    def forward(self, en, de):
        kernels = F.softmax(self.kernel_generator(en, de), dim=1)
        return  self.carafe(de, kernels, self.up_kernel_size, self.scale)


class CARAFE(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, kernel, kernel_size=5, ratio=2):
        B, C, H, W = x.shape
        x = F.unfold(x, kernel_size=kernel_size, stride=1, padding=2)
        x = einops.rearrange(x, 'b (c k_up2) (h w) -> b k_up2 c h w',
                             k_up2=kernel_size**2, w=W)
        x = einops.repeat(x, 'b k c h w -> ratio_2 b k c h w', ratio_2=ratio**2)
        x = einops.rearrange(x, '(r1 r2) b k_up2 c h w -> b k_up2 c (h r1) (w r2)',
                             r1=ratio)
        x = torch.einsum('bkchw,bkhw->bchw',[x, kernel])
        return x
    

if __name__ == '__main__':
    x = torch.randn(2, 3, 4, 16).to('cuda')
    y = torch.randn(2, 3, 8, 32).to('cuda')
    fade = FADE(3).to('cuda')
    print(fade(y, x).shape)

参考文献：

Fusing the assets of encoder and decoder in feature upsampling

SAPA 上一篇

A2U for deep matting 下一篇