Rethinking Semantic Segmentation from a SequencetoSeque
作者：禅与计算机程序设计艺术
1.简介
在图像和视频理解任务中，对像素进行分类是图像分割领域的重要研究对象。其目的是将图像中的物体、人物、场景等进行细化分类。然而，传统的语义分割方法往往受到人类认知能力的限制，无法实现高精度的像素级分类，因此很难用于实际应用。随着计算机视觉领域的飞速发展，许多学者提出了基于深度学习的新型语义分割模型，取得了不错的效果。 Transformer是一种最近被提出的模型，它在自注意力机制上采用了Self-Attention机制，这种机制能够捕获输入序列的信息并且利用信息关联形成新的表示。在语义分割任务中，这种Self-Attention机制可以帮助模型学习到上下文特征之间的关系。因此，通过结合Self-Attention机制和序列到序列（Seq2Seq）模型，Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers便被提出。本文的主要工作包括两方面：一方面是从Seq2Seq模型的角度重新审视语义分割模型；另一方面则是在Seq2Seq的基础上改进语义分割模型。本文的主要贡献如下：
开发了一种创新性的基于序列到序列架构的语义分割模型——Segformer。该模型采用了新型注意力机制，在不同尺度特征间实现了高度的有效融合。该系统经过全面测试，在小样本数据、大规模数据以及模糊不清或多样化场景等多方面的实际应用中均表现出色。
在探索不同特征融合方式及模块设计的过程中, 研究者揭示了对称模块设计在语义分割任务中的重要性. 即使针对相同的任务, 采用不同的模块架构也能够实现着显著的性能差异. 实验结果表明, 在增强学习训练框架下构建的Segformer展现出超越传统模型的优势.
针对多种视觉理解和语言处理任务展开了系统性测试，并成功验证了该模型的有效性
2.基本概念术语说明
首先，我们需要了解一些相关概念及术语。
2.1 Transformer
最近提出的Transformer模型是一种基于注意力机理的深度学习架构。该模型通过自注意力机理捕获输入序列的信息并生成新的表征。该机制能够提取输入序列中各元素及其周围的上下文信息，并根据这些信息间的关联关系对各元素进行处理。基于全局特征融合特性，该机制可有效缓解序列数据维度爆炸的问题。针对语义分割任务，在此框架下可采用Self-Attention机理以探索不同尺度特征间的相互作用关系。
2.2 Seq2Seq模型
该模型属于一种神经网络架构，在机器翻译等任务中具有重要应用价值。它能够将输入的源语言句子转化为相应的目标语言句子。针对语义分割等任务需求，在深度学习领域中已有多个变体被提出并取得不错的效果。该模型主要由编码器和解码器两个核心组件构成：其中编码器负责将输入的空间像素序列作为特征提取的基础；解码器则根据编码结果和当前时间步的信息生成相应的输出特征。值得注意的是，在传统的全连接层设计下无法直接满足这类复杂场景的需求；因此需要结合特定的技术手段进行优化设计。
2.3 Self-Attention机制
自注意力机制是一种关注机制，在模型中帮助提取输入序列中各个元素所包含的上下文信息。
该结构主要由两部分组成：第一部分负责进行线性变换，
第二部分则用于计算各元素间的相似程度并生成对应的权重分配。
通过计算各元素间的相似程度来确定权重分配，
将生成的权重矩阵与上一层输出相乘，
从而得到最终的新表征形式。
这种机制能够有效提取输入序列中的相关信息并据此构建新的表征形式，
并在特定任务场景下应用时展现出独特的优势。
2.4 Multi-Head Attention
Multi-Head Attention是一种基于Self-Attention机制的扩展形式，在实际应用中能够显著增强模型对复杂数据的理解能力。每个注意力头都负责处理特定类型的信息或特征维度，在这种机制下，各个头之间通过加权平均或逐元素相加的方式进行信息整合以生成最终输出结果。在语义分割任务中，这种机制能够有效融合不同尺度的空间信息特征
3.核心算法原理和具体操作步骤以及数学公式讲解
本节将详细阐述本文提出的Segformer模型。该模型的核心创新在于引入了一种新型注意力机制，并采用多头注意力替代传统的卷积核。此外，在模块设计方面，该模型采用了对称架构方案。其优点体现在既能有效减少网络参数规模又能在一定程度上提升其泛化能力。随后我们将依次介绍该模型的主要组件及其工作原理。
3.1 模块设计
该研究旨在增强模型性能，在其架构中整合了四个功能单元并将其分为两类核心组件：上采样分支与下采样路径。在上层结构中（上采样分支），该方法通过自注意力机制与多头结构融合了跨尺度特征信息；而下层路径（下采样路径）则专注于学习并提取各尺度特征间的相互作用关系，并通过1×1卷积进行特征尺度调整以优化表示能力。为减少网络参数规模，在体系结构设计中实现了两个分支结构的有机整合与协同优化，并在此基础上引入多头注意力机制以进一步提升模型性能表现能力。图1展示了该网络架构的整体设计框架及其关键组件配置情况：图1 该网络架构设计框架示意图
3.2 特征处理模块
为应对不同尺寸的输入信号,SemSegmentation形式化模型引入了一个关键组件——'多尺度信息整合器'。该组件能够整合不同尺度的信息并提升其通用适用性,其架构包含多头注意力机制与降采样组件。具体而言,其输出经1×1卷积处理后与原始图像叠加,并叠加降采样后图像;随后,降采样部分通过膨胀卷积核实现降采样并叠加到原始图像上;最终整合后的图像输入至下一层次结构中执行下一步操作
3.3 上采样模块
Segformer主要通过多头注意力机制来进行特征融合，在其设计中特别强调了各子网络间的协同作用特性。具体而言，在上采样模块中，在处理每个特征时仅关注与其直接相关的元素。该模型能够提取出不同尺度特征之间的相互关联关系并加以整合利用，在图2中直观展示了该过程的具体实现流程
3.4 下采样模块
在下采样模块中，Segformer模型识别出不同尺寸特征间的相互联系关系。与其它类似的方法不同的是，在此过程中Segformer特别采用了自注意力机制来建立特征间的联系。该模块采用1×1卷积操作实现特征图尺寸的缩减。为了保证模型性能稳定运行，在此过程中又特意配置了膨胀卷积核来完成下采样过程。如图3所示，该模块主要包含两个关键组件：一是基于自注意力机制的特征融合路径；二是基于膨胀卷积核的特征缩减路径。
3.5 对称的模块设计方案
为了实现降维目标并提升模型性能，Segformer采用了对称模块设计策略。具体而言，在特征处理模块中，该方法通过引入残差连接确保每个注意力层拥有相同数量的参数。在下采样模块中，则采用了对称配置策略来设置两个1x1卷积核以保持完整性。如图4所示展示了这一创新性设计方案
3.6 位置编码
通过向模型注入空间信息的先验知识段落中展示出特定特征属性后段落中详细阐述该技术细节
3.7 训练方法
该模型采用了无监督预训练与微调的联合策略。研究者首先基于ImageNet构建了一个预训练的基础模型。在预训练阶段中，该模型仅完成了第一层自注意力机制参数的学习过程。随后，在后续的微调阶段中，研究者通过一系列多任务学习方法进一步优化了该模型：具体而言，在目标检测任务中主要采用了SSD和YOLOv3算法，在图像分割任务中则主要依赖于FCN和UNet架构。值得注意的是，在目标检测任务中重点进行了边界框回归的学习，在图像分割任务中则着重于实例分割能力的提升。经过微调优化后，该模型在多个推理场景下表现优异：包括目标检测、图像分割等多任务应用能力显著提升。图6详细展示了该模型的整体学习流程框架：包括自监督预训练基础、多模态特征提取以及最终的任务融合优化三个关键步骤。
3.8 小样本学习
为了减小过拟合风险, Segformer 模型采用了小样本学习策略。该模型设定的输入尺寸为 128×128 像素。作者通过随机裁剪的方式获取小于 128×128 像素的小样本数据。在实际应用中,该模型的输入尺寸可灵活设置为较大的尺寸如 224×224 像素。然而,在提升训练效率方面 Batch Normalization 起到了关键作用,Batch Normalization 展现出良好的适应性能力,在不同尺寸特征上表现稳定。由此可见,在处理小样本问题时该模型具有明显优势,这一发现对其他基于全像素建模的方法并不适用。图 7 展示了 Segformer 模型的小样本学习过程
3.9 数据增强
该模型通过数据增强方法实现了对泛化能力的提升。具体而言，该方法采用了多种操作方式：随机裁剪、旋转、翻转、平移以及光亮变化等；其作用在于有效防止过拟合现象的发生。图8详细描绘了该模型的数据增强流程。
3.10 可伸缩性分析
为确保对模型性能进行可靠的评估, 该系统需完成针对特定数据集的测试流程。由此可知, 该方法需进一步优化以提高效率。该系统采用多分支架构设计, 这种架构显著提升了系统的扩展能力, 如图9所示, 其可扩展性特征得到了充分验证。
3.11 模型的性能指标
Segformer模型的各项关键性能数据包括准确性、分割效果、计算效率、参数规模以及计算复杂度等指标。以下将详细阐述Segformer模型的各项关键性能数据。
3.11.1 准确率
该模型在性能上超越了现有所有深度学习模型。该模型展现出卓越的性能，在各种环境下均能表现出色。
3.11.2 分割精度
该模型在不同环境下均表现出较高的分割效果。该研究提出的Segformer方法在标准化评估指标中表现优异。该算法在公开数据集测试中展现出卓越的性能，在PASCAL VOC 2012测试集中达到83.4%的精确度。
3.11.3 推理时间
与现有最新模型相比, Segformer模型在推理速度上的优势明显.
实验结果表明, 在PASCAL VOC2012数据集测试中, 平均单张图片推理时间为30毫秒.
相较于AlexNet而言, 在完成相同任务所需的时间上仅需约一半的时间.
3.11.4 参数数量
Segformer模型的参数规模仅达百万级。因此能够轻松部署在移动端设备上。
3.11.5 FLOPs
Segformer模型的FLOPs比其他最新模型要低得多。
4.具体代码实例和解释说明
本节全面介绍Segformer模型的代码实现过程。本文基于PyTorch开发了该模型，并对其核心算法进行了详细说明。具体提供了Segformer模型的代码实现及其相关说明。
import torch
from torch import nn
class TransformerEncoderLayer(nn.Module):
"""
Implementation of transformer encoder layer for segformer model
Args:
d_model (int): input dimension of feature map
n_heads (int): number of heads in multi head attention layer
dim_feedforward (int): hidden dimension of feed forward network used in self attention and ffn
dropout (float): probability of dropout
"""
def __init__(self, d_model, n_heads, dim_feedforward=2048, dropout=0.1):
super().__init__()
# multi head attention layer
self.attention = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
self.norm1 = nn.LayerNorm(d_model)
# linear layers for feed forward network
self.linear1 = nn.Linear(d_model, dim_feedforward)
self.dropout = nn.Dropout(dropout)
self.linear2 = nn.Linear(dim_feedforward, d_model)
self.norm2 = nn.LayerNorm(d_model)
# dropout layer
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
def forward(self, src, mask=None):
"""
Forward pass of transformer encoder layer
Args:
src (tensor): input tensor of shape [batch size, sequence length, d_model]
mask (bool): boolean mask to zero out non valid elements
Returns:
output (tensor): encoded tensor of same shape as the input tensor
"""
# first multi head attention layer followed by dropout and add residual connection
attn_output, _ = self.attention(src, src, src, attn_mask=mask, need_weights=False)
attn_output = self.dropout1(attn_output)
attn_output = self.norm1(src + attn_output)
# second layer is a fully connected layer followed by activation function (relu or gelu)
# then dropout and add residual connection
intermediate_output = self.linear2(self.dropout(torch.relu(self.linear1(attn_output))))
ffn_output = self.dropout2(intermediate_output)
ffn_output = self.norm2(attn_output + ffn_output)
return ffn_output
class PositionalEncoding(nn.Module):
"""
Implementation of positional encoding for segformer model
Args:
d_model (int): input dimension of feature map
"""
def __init__(self, d_model, max_len=5000):
super().__init__()
pe = torch.zeros(max_len, d_model).float().unsqueeze(0).requires_grad_(False)
position = torch.arange(0, max_len).unsqueeze(1).float()
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.squeeze(0)
self.register_buffer('pe', pe)
def forward(self, x):
"""
Forward pass of positional encoding module
Args:
x (tensor): input tensor of shape [batch size, sequence length, d_model]
Returns:
output (tensor): tensor containing positional information added to the input tensor
"""
x = x + self.pe[:x.size(1)]
return x
class FeatureProcessModule(nn.Module):
"""
Implementation of feature process module for segformer model
Args:
num_patches (int): number of patches to be extracted per image
embed_dims (list): list of embedding dimensions used by each transformer encoder layer
n_heads (list): list of number of heads used by each transformer encoder layer
dropouts (list): list of probabilities of dropout applied after each transformer encoder layer
sr_ratios (list): list of ratios between strides and kernel sizes used by convolutional layers before each transformer encoder layer
mlp_ratio (int): ratio between the hidden dimension of MLP in transformer encoder layer's FeedForward network
"""
def __init__(self, num_patches, embed_dims=[64, 128, 256], n_heads=[1, 2, 4],
dropouts=[0.1, 0.1, 0.1], sr_ratios=[8, 4, 2]):
super().__init__()
patch_embed = PatchEmbedBlock(patch_size=16, embed_dim=embed_dims[0])
self.add_module("patch_embedding", patch_embed)
self.pos_encoding = PositionalEncoding(embed_dims[0])
self.layers = nn.ModuleList([])
for i, (ed, nh, do, srr) in enumerate(zip(embed_dims, n_heads, dropouts, sr_ratios)):
self.layers.append(
TransformerEncoderLayer(
ed,
n_heads=nh,
dim_feedforward=int(mlp_ratio*ed),
dropout=do,
)
)
if srr > 1:
self.add_module(f"sr{srr}_{i}", PatchMergingBlock())
self.norm = nn.LayerNorm(embed_dims[-1])
def forward(self, inputs):
"""
Forward pass of feature processing module
Args:
inputs (tensor): input tensor of shape [batch size, channel, height, width]
Returns:
output (tensor): encoded tensor obtained after passing through all transformer encoder layers
"""
# extract patches using patch embedding block and add positional encoding
x = self.patch_embedding(inputs)
x = self.pos_encoding(x)
for idx, layer in enumerate(self.layers):
# apply transformer encoder layer on the input tensor and concatenate outputs along depth dimension
x = layer(x)
# check if there are any downsampling operations to perform
if isinstance(layer, PatchMergingBlock):
x = getattr(self, f"sr{srr}_{idx}")(x)
# normalize the output vector
x = self.norm(x)
return x
class SegformerModel(nn.Module):
"""
Implementation of Segformer model architecture
Args:
num_classes (int): number of classes to be segmented
backbone (str): name of backbone used by the model (default "resnet50")
pretrained (bool): whether to use pre-trained weights or not
"""
def __init__(self, num_classes, backbone="resnet50", pretrained=True):
super().__init__()
# load resnet50 backbone with optional pre-training on ImageNet dataset
if backbone == "resnet50":
self.backbone = ResNetBackbone(pretrained=pretrained)
else:
raise NotImplementedError("Only ResNet-50 backbone supported currently.")
# set transformer encoder parameters based on resnet50 backbone features
embed_dims = [64, 128, 256, 512]
n_heads = [1, 2, 4, 8]
dropouts = [0.0, 0.0, 0.1, 0.1]
sr_ratios = [8, 4, 2, 1]
mlp_ratio = 4
# create feature process module to convert input images into high level feature maps
self.fpn = FeatureProcessModule(num_patches=(self.backbone.channels//32)**2,
embed_dims=embed_dims, n_heads=n_heads,
dropouts=dropouts, sr_ratios=sr_ratios,
mlp_ratio=mlp_ratio)
# segmentation head to predict final semantic segmentation labels
self.seg_head = SegmentationHead(inplanes=embed_dims[-1], num_classes=num_classes)
def forward(self, inputs):
"""
Forward pass of Segformer model
Args:
inputs (tensor): input tensor of shape [batch size, channel, height, width]
Returns:
output (dict): dictionary containing predicted semantic segmentation masks for different resolutions
"""
# obtain low level features from resnet-50 backbone
feat = self.backbone(inputs)
# convert input images into high level feature maps using feature process module
feat = self.fpn(feat)
# generate semantic segmentation predictions using segmentation head
pred = self.seg_head(feat)
# store predicted segmentation masks for different resolutions
results = {}
_, h, w = pred.shape
for r in np.linspace(0.5, 1.5, 4)[::-1]:
rh, rw = int(h*r+0.5)//32*32, int(w*r+0.5)//32*32
result = F.interpolate(pred, size=(rh,rw), mode='bilinear')[:, :h, :w].contiguous()
results[f"{int(r)}x"] = result
return results
代码解读