RTFormer: Efficient Design for Real-Time Semantic Segmentation with Transformer论文解读

前言

  本文的创新点在于使用了高分辨率与低分辨率并行计算的语义分割transformer框架，并且提出了对低分辨率使用的gpu-frienndly attention，gpu-frienndly attention则是在外部注意力的基础上改进得到的。对高分辨率使用cross attention。
论文地址：论文PDF地址
代码地址：github代码地址

1. stem + basic conv block

  下图为RTFormer的总体架构图。从图中可以得到，共分五个阶段，还要加上开始前的stem阶段和结束后的分割头阶段。
下面我们开始详细介绍。刚开始的stem阶段。实际上是2个步长为2的3x3卷积（k,s,p取3，2，1）。第一个卷积通道数由输入图像的3变为base_channel，默认是64.第二个卷积通道数不变。经过stem阶段后。(n,3,h,w)->(n,base_channel,h/4,w/4).
  接下来是两个类似的阶段，即stage1和stage2.注意这两个阶段还没有分高分辨率和低分辨率。阶段一中特征图主要经过一些卷积和mlp层。具体是过两次basicblock。这两次block后面那个只是单纯的过一遍3x3的ksp取3，1，1的卷积加mlp，也就是说不会对shape造成影响。一个basicblock就是下图的basic conv，包括3x3conv,BN,ReLU，3x3conv，残差结构,ReLU.阶段1的basicblock的ksp取3，1，1.也就是说说stage1的shape没有发生改变。到stage2时channel取2*base_channel，接下来basechannel称为c.stride取2，所以第一个basicblock里的第一个卷积的ksp取3，2，1。在经过stage2后，我们得到了x2，shape为（n,2c,h/8,w/8）。
  在stage3里面和stage2是类似的，通过卷积让通道增加到2倍，h和w各减少到原来的1/2.所以这里得到x3,shape为（n,4c,h/16,w/16),即stage3中位于下方的特征图。这也是低分辨率特征图。将输入x2与经过conv,pooling将C,H,W调整为与x2一致的x3相加。这里获得了高分辨图x3_。

2. RTFormer block

  接下来是比较重要的RTFormer块。也就是上图的红色模块。下图是RTFormer块的内部结构。

  我们可以看到现在是高分辨率与低分辨率特征图并行计算了。本文对它们采用的attention是不同的。对低分辨率图像，使用GPU-friendly attention,对高分辨率图像，采用cross-attention。下面分别介绍这两种attention。图中出现的残差结构将不再讲述，残差结构往往通过conv,pooling(或者是插值interpolate)将C,H,W调整为与和目标图一致，然后相加。
例如x3_=x2+pooling(conv(x2))：

self.compression3 = nn.Sequential(
            bn2d(base_chs * 4),
            nn.ReLU(),
            conv2d(
                base_chs * 4, base_chs * 2, kernel_size=1), )
x3_ = x2 + F.interpolate(
            self.compression3(x3), size=paddle.shape(x2)[2:], mode='bilinear')

  先说GPU-friendly attention ,它本质上是对external attention的修改版本。实质上就是外部注意力把对qk^T与（qk^T）v的两次计算使用线性层实现,相当于输入特征图的Q外部KV做了计算。（过一遍线性层相当于与一个矩阵相乘）
在本文中，GPU-friendly attention，下称GFA，在外部注意力的基础上做了一些改变。包括不再做多头运算，因为文中认为多头需要对矩阵进行拆分，这对GPU十分不友好。所以GFA放弃了多头计算。但是为了保留多头计算的好处，对计算保留了在N=12²进行激活和归一化。因为指定的外部K的shape

可以看到，k,v都是可学习的参数。实现外部注意力的方式被换成了F.conv2d。将k指定为卷积核。qkt将k指定为卷积核对q做卷积，k和q的shape与典型的attention是不一样的。cross_k的shape是（n144,c_in,1,1）因为他要做卷积核，所以（1，1）代表1x1卷积,c_in是in_channel，n144代表中间维度。在做GFA的时候，因为qk^Tstride取2所以H,W为1/32.

	def __init__():
        if use_cross_kv:
            assert self.same_in_out_chs, "in_channels is not equal to out_channels when use_cross_kv is True"
        else:
            self.k = self.create_parameter(
                shape=(inter_channels, in_channels, 1, 1),
                default_initializer=paddle.nn.initializer.Normal(std=0.001))
            self.v = self.create_parameter(
                shape=(out_channels, inter_channels, 1, 1),
                default_initializer=paddle.nn.initializer.Normal(std=0.001))
    def forward(self, x, cross_k=None, cross_v=None):
        """
        Args:
            x (Tensor): The input tensor. 
            cross_k (Tensor, optional): The dims is (n*144, c_in, 1, 1)#cross_k做卷积核。c_in是in_CHANNEL,n*144是out_chammel
            cross_v (Tensor, optional): The dims is (n*c_in, 144, 1, 1)
        """
        x = self.norm(x)
        if not self.use_cross_kv:
            x = F.conv2d(#qkt
                x,
                self.k,
                bias=None,
                stride=2 if not self.same_in_out_chs else 1,
                padding=0)  # n,c_in,h,w -> n,c_inter,h,w
            x = self._act_dn(x)  # n,c_inter,h,w
            x = F.conv2d(#(qkt)v
                x, self.v, bias=None, stride=1,
                padding=0)  # n,c_inter,h,w -> n,c_out,h,w
        else:
            assert (cross_k is not None) and (cross_v is not None), \
                "cross_k and cross_v should no be None when use_cross_kv"
            B = x.shape[0]
            assert B > 0, "The first dim of x ({}) should be greater than 0, please set input_shape for export.py".format(
                B)
            x = x.reshape([1, -1, 0, 0])  # n,c_in,h,w -> 1,n*c_in,h,w
            x = F.conv2d(
                x, cross_k, bias=None, stride=1, padding=0,#cross_k:n*144,out_channels_h,1,1
                groups=B)  # 1,n*c_in,h,w -> 1,n*144,h,w  (group=B)
            x = self._act_sn(x)
            x = F.conv2d(
                x, cross_v, bias=None, stride=1, padding=0,#cross_v:n*out_channels_h,144,1,1
                groups=B)  # 1,n*144,h,w -> 1, n*c_in,h,w  (group=B)
            x = x.reshape([-1, self.in_channels, 0,
                           0])  # 1, n*c_in,h,w -> n,c_in,h,w  (c_in = c_out)
        return x

  cross-attention代码见上图

3. DAPPM

将多尺度特征图进行融合以便于得到分割结果。
将输入的特征图通过步长为2,4,8,16的nn.AvgPool2D函数与一系列的BN,ReLU,1x1conv后变成长宽不同的多尺度特征图，然后将其通过累加前项的方法来获得x1到x4,x0没有进行池化。只有一系列MLP层。最后将多尺度特征图concat再通过conv让通道数回归out_channel（2c），再通过seghead可以让通道数变为分类的数量。这样就生成了我们需要的分割图。