深度探讨CrossFormer如何解决跨尺度问题

AIWalker

发布于 2021-08-25 11:27:44

1.1K00

代码可运行

文章被收录于专栏：AIWalkerAIWalker

运行总次数：0

代码可运行

极市导读

浙大联合腾讯等开源的新视觉模块CrossFormer最近开源，该工作通过提出两个模块：CEL和LSDA，弥补了以往架构在建立跨尺度注意力方面的缺陷。本文作者对其进行了详细的分析，介绍了各模块的设计原理模型结构，更深入的了解工作的核心。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

论文链接: https://arxiv.org/pdf/2108.00154.pdf

论文代码：https://github.com/cheerss/CrossFormer

1. 出发点

Transformers模型在处理视觉任务方面已经取得了很大的进展。然而，现有的vision transformers仍然不具备一种对视觉输入很重要的能力：在不同尺度的特征之间建立注意力。

每层的输入嵌入都是等比例的，没有跨尺度的特征；
一些transformers模型为了减少self-attention的计算量，衰减了key和value的部分特征表达。

2. 怎么做

为了解决上面的问题，提出了几个模块。

Cross-scale Embedding Layer (CEL)
Long Short Distance Attention (LSDA)
Dynamic Position Bias (DPB)

这里1和2都是为了弥补了以往架构在建立跨尺度注意力方面的缺陷，3的话和上面的问题无关，是为了使相对位置偏差更加灵活，更好的适合不定尺寸的图像和窗口。这篇文章还挺讲究，不仅提出两个模块来解决跨尺度特征attention，还附送了一个模块来搞一个搞位置编码。

3. 模型结构

模型结构

模型整体的结构图如上所示，与swin-transformers和pvt基本整体结构一致，都是采用了层级的结构，这样的好处是可以迁移到dense任务上去，做检测，分割等。整体结构由以下组成：

Cross-scale embeeding layer (CEL) , 用来做patch embeeding和patch merging(下采样)。
CrossFrom block，看上图(b)，整体看是两个transformer结构的block所组成，其中第一个transformer block采用的是SDA，也就是short distance attention,并且引入了一个DPB模块，第二个transformer block采用的则是LDA，也就是long distance attention，同样也引入了一个DPB模块，两个transformer block串行，组成一个CrossFormer block。
Classification Head，就是常规的分类MLP，没啥可说的。

3.1 Cross-scale embeeding layer (CEL)

Q&A

Question:既然是层级结构，那么就一定会有尺度上的下采样，那crossformers是怎么做的呢？

Answer: 简单回顾一下pvt和swin的做法

pvt: 假设feature map为, 那么我们就可以做一个stride为2的一个convolution, 变换为，由于patchsize固定，所以，featuremap下采样，对应的就是token的下采样。

swin: swin由于是基于windows做attention，为了达到下采样的效果，选择直接对featuremap上采样，每个4邻域都会分别采样到另一个map里面去，最后则有变换为，也可以看做是stride为2带有空洞的卷积操作。

Question: 万变不离其宗，所以为了达到下采样的效果，用卷积其实就可以了。那么CrossFormer为了实现下采样是怎么做的呢？

Answer:

patch embeeding

看上图，很明显，直接用不同卷积核来对输入的图片做卷积，得到卷积后的结果，直接concat一起，作为我们的patch embeeding。想法很简单，实现的话也很朴素，通过不同卷积核的卷积，来获取不同尺度特征的信息，对于变化尺度的物体相对来说是比较友好的，这个可行性其实在很多paper里面都有用到过，比如Pyramidal Convolution, 如下图所示。

尺度卷积

ps: 这里除了patch embeeding，也就是第一个CEL用的是4个卷积核stride为4来做多尺度，其余的CEL也就是patch merge用的都是2个卷积核stride为2来做的多尺度。两个操作基本相同，只看一份代码即可，核心代码如下:

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=[4], in_chans=3, embed_dim=96, norm_layer=None):
        super().__init__()
        ...

        self.projs = nn.ModuleList()
        for i, ps in enumerate(patch_size):
            if i == len(patch_size) - 1:
                dim = embed_dim // 2 ** i
            else:
                dim = embed_dim // 2 ** (i + 1)
            stride = patch_size[0]
            padding = (ps - patch_size[0]) // 2
            self.projs.append(nn.Conv2d(in_chans, dim, kernel_size=ps, stride=stride, padding=padding))
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        xs = []
        for i in range(len(self.projs)):
            tx = self.projs[i](x).flatten(2).transpose(1, 2)
            xs.append(tx)  # B Ph*Pw C
        x = torch.cat(xs, dim=2)
        if self.norm is not None:
            x = self.norm(x)
        return x

代码做了两件事情:

初始化几个不同kernel，不同padding，相同stride的conv
对输入进行卷积操作后得到的feature，做concat

这样, 以输入为224x224为例, 我们通过patch embeeding, 得到了一个56x56的featuremap，输入到第一个stage，输出继续做一个patchmerging，得到了一个28x28的featuremap，输入到第二个stage，输出继续做一个patchmerging，得到了一个14x14的featuremap, 输入到第三个stage, 输出再次做一个patchmerging，得到一个7x7的featuremap，在输入到最后一个stage，最后的输出做分类即可，基本上都是这么一个套路了，大同小异。那么stage里面是怎么做的，看下一节。

3.2 Stage block

对于标准的transformerblock来说，假设输入为, 经过transformer后，我们的输出还是，输入和输出是没有变化的，唯一的尺度变换都在patch embeeding和patch merging。那么我们在改动transformer block的时候，也是要遵守这一原则，对应的，如果想有resolution上的变化，那么就要借助于reshape或者view等操作，好了，不说废话，看这篇文章的crossformer block是怎样的。

CrossFormer Block

cfblock

CrossFormer block由两个transformer的block堆叠而成，两个transformer block的self-attention都是基于windows来做的，不同之处在于一个考虑的是局部内的信息，一个则是考虑的是全局的信息。这个思想并没有什么突出的地方，目前来说transformer做局部和全局的串联，已经屡见不鲜。

Q&A

Question: 问题来了，怎样实现呢，既要保证基于windows做self-attention，又想要全局的信息？

Answer: 使用一个固定的步长step，比如2或者3，对行和列分别按步长采样，这样可以得到多个全局的信息，同时基于一个大小的windows。这样最大可能的利用到了featuremap的全局性，同时节省了计算的复杂度，假设输入为，step为，那么windows的窗口大小为，原始的复杂度为, 那么基于窗口的attention的复杂度为。

CrossFormerblock中的基石: windows self-attention

Short Distance Attention(SDA)

SDA

对于一个的feautremap，如果我们想要实现self-attention, 需要先转换为的向量，那么这里就是所谓的long-range的attention，也就是全局的。但是对于MHA来说，部分head还是更多的focus到short-range，结合swin和twins的结论可以验证，局部attention不仅可以达到很好的效果同时还会节省计算。那么怎么获得局部的attention，很简单，如上图所示，只需要把原始的做reshape操作, 既可以得到，那么我们只需要对4个做attention即可，最后在reshape回原始形状，代码如下:

x = x.reshape(B, H // G, G, W // G, G, C).permute(0, 1, 3, 2, 4, 5)
x = attention(x)
....

Long Distance Attention(LDA)

LDA

从上面的SDA, 我们得到了局部attention，但是也说了，部分head是局部友好的，也就是说，对于self-attention来说，long-range始终是必不可少的，所以还是需要引入long distance attention。如上图所示，颜色一致的部分表示的是归属于同一个sub-windows的，对于原始的，使用step为2进行采样，得到了4个, 可以抽象成两种计算方法，一种是空洞卷积，一种则是1x1的卷积，stride为step，对于图像来说，相邻的位置，像素所表达的信息接近，所以两种得到的都是全局的一个感受野，所以对应我们的attention，也会得到一个近乎全局的attention，代码如下:

x = x.reshape(B, G, H // G, G, W // G, C).permute(0, 2, 4, 1, 3, 5)
x = attention(x)
...

直接看这个代码可能不太好理解，我们用einops简单改写一下，代码如下: 输入：

x[0,:,:,0]
tensor([[ 1,  2,  3,  4],
        [ 2,  4,  6,  8],
        [ 3,  6,  9, 12],
        [ 4,  8, 12, 16]])
x.shape
torch.Size([1, 4, 4, 1])

SDA:

a1 = rearrange(x, ' b (h g1) (w g2) c -> b h w g1 g2 c ', g1=2, g2=2)
a1[0,:,:,:,:,0]
tensor([[[[ 1,  2],
          [ 2,  4]],

         [[ 3,  4],
          [ 6,  8]]],


        [[[ 3,  6],
          [ 4,  8]],

         [[ 9, 12],
          [12, 16]]]])

对于SDA的情况，实际上就是循环HW，扣2x2的区域下来，那么因为有行遍历优先，或者列遍历优先，实际上得到的结果是顺序的。

LDA:

a2 = rearrange(x, ' b (g1 h) (g2 w) c -> b h w g1 g2 c ', g1=2, g2=2)
a2[0,:,:,:,:,0]
tensor([[[[ 1,  3],
          [ 3,  9]],

         [[ 2,  4],
          [ 6, 12]]],


        [[[ 2,  6],
          [ 4, 12]],

         [[ 4,  8],
          [ 8, 16]]]])

那么对于LDA的情况，我们希望的是外循环是有间隔的，所以把step放到HW的外面，这样循环的时候则是按间隔来进行sample，以达到全局的效果。

CrossFormerblock中的位置编码: Dynamic Position Bias(DPB)

Relative position bias (RPB)随着位置编码技术的不断发展，相对位置编码偏差逐渐的应用到了transformers中，很多的vision transformers均采用RPB来替换原始的APE，好处是可以直接插入到我们的attention中，不需要很繁琐的公式计算，并且可学习性高，鲁棒性强，公式如下:

Attention = Softmax(QK^{T}/ \sqrt{d} + B)V

Q&A

Question:但是这里有个问题，对于来说，会有一个偏差, 所表达的则是matrix上的i和j的相对位置的embeeding，很显然，如果图像的尺寸变化，那么可能会超出B所表达的范围，会导致PE没有作用，那么要怎么改进呢？

Answer:很简单，插值或者切片不就好了，但是切片会导致pe完整性差，损失信息，插值是通过原始的位置信息来模拟出来信息，实际上还是原始的信息，没有信息收益。那本文想到的一个方法就是可以通过学习得到位置信息。

Dynamic Position Bias (DPB)

DPB

举个栗子，如果我们的窗口大小为, 那么我们希望的相对位置范围为假设我们不考虑截断距离，如果我们的窗口突然放大到了，那么我们实际的相对位置所表达的信息只是中间的一部分窗口，失去了对外层数据位置的访问。DPB的思想则是，我们不希望通过用实际的相对位置来做embeeidng，而是希望通过隐空间先对位置偏差进行学习，如上图所示。

DPB，由3个线性层+LayerNorm+ReLU组成的block堆叠而成，最后接一个输出为1的线性层做bias的表征，输入是，由于self-attention是由多个head组成的，所以输出为，代码如下：

先得到一个相对位置偏差的矩阵，假设group_size的大小为，那么bias的维度为 self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False) # generate mother-set position_bias_h = torch.arange(1 - self.group_size[0], self.group_size[0]) position_bias_w = torch.arange(1 - self.group_size[1], self.group_size[1]) biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w])) # 2, 2Wh-1, 2Wh-1 biases = biases.flatten(1).transpose(0, 1).float() self.register_buffer("biases", biases) biases: tensor([[-6., -6.], [-6., -5.], [-6., -4.], [-6., -3.], [-6., -2.], [-6., -1.], ... [ 6., 4.], [ 6., 5.], [ 6., 6.]])
初始化DBP模块 pos = DynamicPosBias(64 // 4, 8, residual=False)

Rethinking: 对于PE来说，目前的形成方法都是通过embeeding来构建bias矩阵，对于VIT来说，直接使用绝对位置的embeeding，通过学习来更新，对于swins来说，直接使用embeeding而不是相对bias的值，相当于，其实本质上没有太大的差异, 从消融实验结果上来看，DBP和RBP的性能一样。唯一的作用，就是embeeding是后验而不是先验，对于变换的尺寸来说，可能更加友好，只不过这个paper里面没有给出结论，还需要更多的实验来验证。

DBP&RBP

综上，我们每个stageblock里面，都是由SDA+DBP&LDA+DBP堆叠而成，与swin类似，奇数layer走SDAblock，偶数layer走LDAblock，从结构上来看，先局部attention，再全局attention，有一点点由点到面的既视感。

模型设计

与其他的paper大同小异了，设计了4种不同FLOPs的模型，Tiny, Small, Big和Large 用来和其他的模型在同等FLOPs下公平比较。表示的是维度，表示的是attention头的个数，表示的是attention窗口的大小，表示的是滑动窗口的间隔。

4. 实验结果

imagenet

CrossFormer都是再224x224的图片大小下进行训练，使用的类似DeiT的训练策略，不过采用了更大的warmup(20个，DeiT是5), 学习率为1e-3, weightdecay为5e-2, 与DeiT不同的是，这里随着模型大小的改变，分别采用了0.1，0.2，0.3，0.5的drop path rate。可以看到，在同等数量级的FLOPs的情况下，CF在imagenet上都取得了SOTA的效果。

detection&segmentation

可以看到CrossFromer在coco2017上基于RetinaNet架构，也可以达到SOTA的效果，高于Twins模型1.4个ap之多。实例分割则是基于Mask-Rcnn的架构，也是SOTA，超过Swin 1.7个ap。相比而言参数量和FLOPs都更少，性能更好。

segmentation

语义分割上，可以看到可以看到最多提升3.3%的MIOU，非常厉害了。

module

消融实验上，可以看到，当CEL和LSDA一起使用的时候，性能最高。不过这实验也很明显了，CrossFormer参考了PVT和swin的设计思想。使用了LSDA，相比于Swin提升了0.6%个点，设计比swin更加朴实，不错的提升。

5. 结论

本文提出了一个新的transformers架构称为CrossFormer。其核心设计包括一个跨尺度嵌入层（CEL）和长短距离注意（LSDA）模块。此外，我们提出了动态位置偏置（DPB），使相对位置偏置适用于任何输入尺寸。实验表明，CrossFormer在几个有代表性的视觉任务上取得了SOTA。特别是，CrossFormer在检测和分割方面有很大的改进，这表明跨尺度嵌入和LSDA对于密集预测的视觉任务特别重要。

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2021-08-09，如有侵权请联系 cloudcommunity@tencent.com 删除

css