Segformer论文研读
论文详细讲解如下：
【论文笔记】Segformer论文阅读笔记_嘟嘟太菜了的博客-CSDN博客
论文笔记——Segformer: 一种基于Transformer的语义分割方法 - 知乎
上图上segformer的整体结构
----------------------------------------------------------------------------------------------------------------------
segformer = backbone(mit) + segformerHead
1、backbone(mit):
import torch
from torch import nn, Tensor
from torch.nn import functional as F
from semseg.models.layers import DropPath
"""
Attention模块计算注意力
例如:
输入：x:[1,3136,32]
输出：x:[1,3136,32]
"""
class Attention(nn.Module):
# dim=32,head=1,sr_ratio=8
def __init__(self, dim, head, sr_ratio):
super().__init__()
self.head = head
self.sr_ratio = sr_ratio # self.sr_ratio = 8
self.scale = (dim // head) ** -0.5 # self.scale 用于将q * K的转置的结果 * self.scale
self.q = nn.Linear(dim, dim)
self.kv = nn.Linear(dim, dim*2)
self.proj = nn.Linear(dim, dim)
if sr_ratio > 1:
self.sr = nn.Conv2d(dim, dim, sr_ratio, sr_ratio)
self.norm = nn.LayerNorm(dim)
# x:[1,3136,32] 输入
def forward(self, x: Tensor, H, W) -> Tensor:
B, N, C = x.shape # B 为batch_size ,N为每个patch的长*宽，C为通道数
# q:[1,3136,32] q的维度和输入x保持一致
q = self.q(x)
# q:[1,1,3136,32]
q = q.reshape(B, N, self.head, C // self.head).permute(0, 2, 1, 3)
if self.sr_ratio > 1:
# x:[1,32,56,56]
x = x.permute(0, 2, 1).reshape(B, C, H, W)
# x:[1,32,7,7]
x = self.sr(x)
# x:[1,49,32]
x = x.reshape(B, C, -1).permute(0, 2, 1)
# x:[1,49,32]
x = self.norm(x)
# kv:[1,49,64]
kv = self.kv(x)
# k:[1,1,49,32]
# v:[1,1,49,32]
k, v = kv.reshape(B, -1, 2, self.head, C // self.head).permute(2, 0, 3, 1, 4)
# @矩阵相乘运算符
# attn:[1,1,3136,49]
attn = (q @ k.transpose(-2, -1)) * self.scale
# attn:[1,1,3136,49]
attn = attn.softmax(dim=-1)
# x:[1,3136,32]
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
# x:[1,3136,32]
x = self.proj(x)
return x
"""
深度可分离卷积
输入：[1,3136,128]
输出：[1,3136,32]
"""
class DWConv(nn.Module):
def __init__(self, dim):
super().__init__()
self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, groups=dim)
# x:[1,3136,128]
def forward(self, x: Tensor, H, W) -> Tensor:
B, _, C = x.shape
x = x.transpose(1, 2).view(B, C, H, W) # x:[1,128,56,56]
x = self.dwconv(x) # x:[1,128,56,56]
return x.flatten(2).transpose(1, 2)
"""
MLP多层感知机
输入:[1,3136,32]
输出:[1,3136,32]
"""
class MLP(nn.Module):
def __init__(self, c1, c2):
super().__init__()
self.fc1 = nn.Linear(c1, c2)
self.dwconv = DWConv(c2)
self.fc2 = nn.Linear(c2, c1)
# def forward(self, x: Tensor, H, W) -> Tensor:
# return self.fc2(F.gelu(self.dwconv(self.fc1(x), H, W)))
# x:[1,3136,32]
def forward(self, x: Tensor, H, W) -> Tensor:
x = self.fc1(x) # x:[1,3136,128]
x = self.dwconv(x, H, W) # [1,3136,32]
x = self.fc2(F.gelu(x))
return x
"""
功能：通过conv操作将图像划分为不同的patch,再把维度做一个变换
例如：
输入：[1,3,224,224]
过程：[1，3,224,224] --conv--> [1,32,56,56] --flatten,transpose--> [1,3136,32],其中3136=56*56
输出：[1,3136,32]
"""
class PatchEmbed(nn.Module):
def __init__(self, c1=3, c2=32, patch_size=7, stride=4):
super().__init__()
self.proj = nn.Conv2d(c1, c2, patch_size, stride, patch_size//2) # padding=(ps[0]//2, ps[1]//2)
self.norm = nn.LayerNorm(c2)
# x:[1,3,224,224]
def forward(self, x: Tensor) -> Tensor:
# x:[1,32,56,56]
x = self.proj(x)
_, _, H, W = x.shape
# x:[1,3136,32]
x = x.flatten(2).transpose(1, 2)
# x:[1,3136,32]
x = self.norm(x)
return x, H, W
"""
Block为Mit模型的基本结构
输入：[1,3136,32]
输出：[1,3136,32]
"""
class Block(nn.Module):
def __init__(self, dim, head, sr_ratio=1, dpr=0.):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = Attention(dim, head, sr_ratio)
self.drop_path = DropPath(dpr) if dpr > 0. else nn.Identity()
self.norm2 = nn.LayerNorm(dim)
self.mlp = MLP(dim, int(dim*4))
# def forward(self, x: Tensor, H, W) -> Tensor:
# x = x + self.drop_path(self.attn(self.norm1(x), H, W))
# x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
# return x
# x：[1,3136,32]
def forward(self, x: Tensor, H, W) -> Tensor:
# Efficient Self-Atten 模块
x1 = self.norm1(x) # x1:[1,3136,32]
x1 = self.attn(x1, H, W) # x1:[1,3136,32] 此模块计算注意力
x1 = self.drop_path(x1) # x1:[1,3136,32] drop_path作用：若x为输入的张量，其通道为[B,C,H,W],那么drop_path的含义为在一个batch_size中，随机有dpr的样本，不经过主干，而直接由分支进行恒等映射
x = x + x1
# Mix-FFN模块
x2 = self.norm2(x)
x2 = self.mlp(x2, H, W) # x2：[1,3136,32] 多层感知机
x2 = self.drop_path(x2)
x = x + x2
return x
mit_settings = {
'B0': [[32, 64, 160, 256], [2, 2, 2, 2]], # [embed_dims, depths]
'B1': [[64, 128, 320, 512], [2, 2, 2, 2]],
'B2': [[64, 128, 320, 512], [3, 4, 6, 3]],
'B3': [[64, 128, 320, 512], [3, 4, 18, 3]],
'B4': [[64, 128, 320, 512], [3, 8, 27, 3]],
'B5': [[64, 128, 320, 512], [3, 6, 40, 3]]
}
"""
Mit为SegFormer的Encoder部分，主要用于特征提取；
输入：[1,3,224,224]
输出：[1,32,56,56],[1,64,28,28],[1,160,14,14],[1,256,7,7] 多尺度信息
"""
class MiT(nn.Module):
def __init__(self, model_name: str = 'B0'):
super().__init__()
assert model_name in mit_settings.keys(), f"MiT model name should be in {list(mit_settings.keys())}" # model name 必须在mit_settings中
embed_dims, depths = mit_settings[model_name] # model name:B0， depths:[2,2,2,2],embed_dims:[32,64,160,256]
drop_path_rate = 0.1
self.channels = embed_dims # self.channels：[32,64,160,256]
# patch_embed
self.patch_embed1 = PatchEmbed(3, embed_dims[0], 7, 4)
self.patch_embed2 = PatchEmbed(embed_dims[0], embed_dims[1], 3, 2)
self.patch_embed3 = PatchEmbed(embed_dims[1], embed_dims[2], 3, 2)
self.patch_embed4 = PatchEmbed(embed_dims[2], embed_dims[3], 3, 2)
# 将drop_path_rate等间隔分为sum(depths),即8份
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
cur = 0
# depths[0] = 2，即第一层包含2个block的深度,第一个block(32,1,8,dpr[0]),第二个block（32,1,8,dpr[1]）
self.block1 = nn.ModuleList([Block(embed_dims[0], 1, 8, dpr[cur+i]) for i in range(depths[0])])
self.norm1 = nn.LayerNorm(embed_dims[0])
cur += depths[0]
self.block2 = nn.ModuleList([Block(embed_dims[1], 2, 4, dpr[cur+i]) for i in range(depths[1])])
self.norm2 = nn.LayerNorm(embed_dims[1])
cur += depths[1]
self.block3 = nn.ModuleList([Block(embed_dims[2], 5, 2, dpr[cur+i]) for i in range(depths[2])])
self.norm3 = nn.LayerNorm(embed_dims[2])
cur += depths[2]
self.block4 = nn.ModuleList([Block(embed_dims[3], 8, 1, dpr[cur+i]) for i in range(depths[3])])
self.norm4 = nn.LayerNorm(embed_dims[3])
# x:[1,3,224,224]
def forward(self, x: Tensor) -> Tensor:
B = x.shape[0]
# stage 1
# x:[1,3136,32],H:56,W:56
x, H, W = self.patch_embed1(x)
for blk in self.block1:
x = blk(x, H, W)
# x:[1,3136,32]
# x1:[1,32,56,56]
x1 = self.norm1(x).reshape(B, H, W, -1).permute(0, 3, 1, 2)
# stage 2
# x:[1,784,64],H:28,W:28
x, H, W = self.patch_embed2(x1)
for blk in self.block2:
x = blk(x, H, W)
# x:[1,784,64]
# x2:[1,64,28,28]
x2 = self.norm2(x).reshape(B, H, W, -1).permute(0, 3, 1, 2)
# stage 3
# x:[1,196,160],H:14,W:14
x, H, W = self.patch_embed3(x2)
for blk in self.block3:
x = blk(x, H, W)
# x:[1,196,160]
# x3:[1,160,14,14]
x3 = self.norm3(x).reshape(B, H, W, -1).permute(0, 3, 1, 2)
# stage 4
# x3:[1,160,14,14]
# x:[1,49,256]
x, H, W = self.patch_embed4(x3)
for blk in self.block4:
x = blk(x, H, W)
# x:[1,49,256]
# x4:[1,256,7,7]
x4 = self.norm4(x).reshape(B, H, W, -1).permute(0, 3, 1, 2)
return x1, x2, x3, x4
if __name__ == '__main__':
model = MiT('B0')
x = torch.zeros(1, 3, 224, 224)
outs = model(x)
for y in outs:
print(y.shape)
2、segformerHead:
import torch
from torch import nn, Tensor
from typing import Tuple
from torch.nn import functional as F
class MLP(nn.Module):
def __init__(self, dim, embed_dim):
super().__init__()
self.proj = nn.Linear(dim, embed_dim)
def forward(self, x: Tensor) -> Tensor:
x = x.flatten(2).transpose(1, 2)
x = self.proj(x)
return x
class ConvModule(nn.Module):
def __init__(self, c1, c2):
super().__init__()
self.conv = nn.Conv2d(c1, c2, 1, bias=False)
self.bn = nn.BatchNorm2d(c2) # use SyncBN in original
self.activate = nn.ReLU(True)
def forward(self, x: Tensor) -> Tensor:
return self.activate(self.bn(self.conv(x)))
class SegFormerHead(nn.Module):
def __init__(self, dims: list, embed_dim: int = 256, num_classes: int = 19):
super().__init__()
for i, dim in enumerate(dims):
self.add_module(f"linear_c{i+1}", MLP(dim, embed_dim))
self.linear_fuse = ConvModule(embed_dim*4, embed_dim)
self.linear_pred = nn.Conv2d(embed_dim, num_classes, 1)
self.dropout = nn.Dropout2d(0.1)
def forward(self, features: Tuple[Tensor, Tensor, Tensor, Tensor]) -> Tensor:
B, _, H, W = features[0].shape
outs = [self.linear_c1(features[0]).permute(0, 2, 1).reshape(B, -1, *features[0].shape[-2:])]
for i, feature in enumerate(features[1:]):
cf = eval(f"self.linear_c{i+2}")(feature).permute(0, 2, 1).reshape(B, -1, *feature.shape[-2:])
outs.append(F.interpolate(cf, size=(H, W), mode='bilinear', align_corners=False))
seg = self.linear_fuse(torch.cat(outs[::-1], dim=1))
seg = self.linear_pred(self.dropout(seg))
return
if __name__ == "__main__":
model = SegFormerHead([32,64,160,256],256,19)
x1 = torch.zeros([1,32,56,56])
x2 = torch.zeros([1,64,28,28])
x3 = torch.zeros([1,160,14,14])
x4 = torch.zeros([1,256,7,7])
features = (x1,x2,x3,x4)
outs = model(features)
print(outs.shape)3、segformer:https://github.com/sithu31296/semantic-segmentation/blob/main/semseg/models/segformer.py
import torch
from torch import Tensor
from torch.nn import functional as F
from semseg.models.base import BaseModel
from semseg.models.heads import SegFormerHead
class SegFormer(BaseModel):
def __init__(self, backbone: str = 'MiT-B0', num_classes: int = 19) -> None:
super().__init__(backbone, num_classes)
self.decode_head = SegFormerHead(self.backbone.channels, 256 if 'B0' in backbone or 'B1' in backbone else 768, num_classes)
self.apply(self._init_weights)
def forward(self, x: Tensor) -> Tensor:
y = self.backbone(x)
y = self.decode_head(y) # 4x reduction in image size
y = F.interpolate(y, size=x.shape[2:], mode='bilinear', align_corners=False) # to original image shape
return y
if __name__ == '__main__':
model = SegFormer('MiT-B0')
# model.load_state_dict(torch.load('checkpoints/pretrained/segformer/segformer.b0.ade.pth', map_location='cpu'))
x = torch.zeros(1, 3, 512, 512)
y = model(x)
print(y.shape)