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

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块的内部结构。
rtformer block
  我们可以看到现在是高分辨率与低分辨率特征图并行计算了。本文对它们采用的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的修改版本。实质上就是外部注意力把对qkT与(qkT)v的两次计算使用线性层实现,相当于输入特征图的Q外部KV做了计算。(过一遍线性层相当于与一个矩阵相乘)
在本文中,GPU-friendly attention,下称GFA,在外部注意力的基础上做了一些改变。包括不再做多头运算,因为文中认为多头需要对矩阵进行拆分,这对GPU十分不友好。所以GFA放弃了多头计算。但是为了保留多头计算的好处,对计算保留了在N=122进行激活和归一化。因为指定的外部K的shape
多头外部注意力与GPU友好型注意力
可以看到,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的时候,因为qkTstride取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可以让通道数变为分类的数量。这样就生成了我们需要的分割图。
DAPPM

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值