参考代码:SegFormer
1.概述
介绍:这篇文章提出的分割方法是基于transformer结构构建的,不过这里使用到的transformer是针对分割任务在patch merge、self-attention和FFN进行了改进,使其更加适合分割任务(无需position-encoding,测试图片的尺寸带来的影响更小)对上下文语义信息和局部细节信息的需求并且更加整体模型轻量化,同时得益于transformer强大的上下文感知能力使得对于解码器只使用简单的几层全连接便可达到更好的分割性能。通观全文最大的亮点在分割网络编码器的设计上,若是将对应的解码器再进行进一步调试分割的效果将更好,不过使用如ASPP这样的感受野增强模块效果就不是那么大了。
将文章的方法于其它的一些分割方法进行比较,如下图。文章方法的分割性能更好且参数量更少。
2. 方法设计
2.1 pipeline
文章方法整体上还是一个编解码网络结构,使用新设计的编码器MiT-Bx系列和多层全连接组成,其结构见下图所示:
在上图中最重要的部分便是其中的transformer编码器了,不过这里针对分割任务设定了如下要求:
- 1)能输出多尺度的编码器特征,使得分割网络可以有效使用上下文语义信息和物体细节信息。对应在Swin Transformer中是通过类似于pixle-shuffle反向操作的形式实现多尺度输出的(可以参考这篇文章的解读:图解Swin Transformer);
- 2)网络设计需要轻量化,本身分割任务就是一种稠密预测任务,轻量化是在实际部署过程中无法回避的问题;
- 3)原本的ViT中存在position-encoding这个在固定输入尺寸的情况下问题不大,但是对于分割这种对于分辨率影响比较大的任务,训练size和infer size不一样会导致position encoding存在问题,因而需要对position-encoding这块进行改进;
2.2 Transformer Block
该模块是文章编码器的基础组成模块,其中包含了:高效self-attention、Mix-FFN和Patch Merge模块,其结构见下图所示:
PS: 这里结构的顺序和文章代码中提供的顺序不太一致,代码中是patch merge在最前面。
patch merge:
在原本的ViT中是单独对每个patch独立进行处理,但是对于分割这类任务来说每个patch单独进行处理且之间无关联是不符合实际情况的,因而需要建立起不同patch之间的联系。在Swin Transformer中是通过Shifted Window Attention来实现不同window上信息的交互的,而在这篇文章中是通过卷积的形式,也就是将每个patch看作是一个卷积窗口,并以这个窗口在输入数据上进行滑动。同时还通过控制其中的stride参数控制输出特征图的大小,从而实现多尺度特征图的输出(卷积带padding)。
class OverlapPatchEmbed(nn.Module):
""" Image to Patch Embedding
"""
def __init__(self, img_size=224, patch_size&#