【悉读经典】SegFormer:语义分割中的层次化Transformer网络

# patch_embed，通过定义卷积操作的步长/stride，时相下采样self.patch_embed1 = OverlapPatchEmbed(    img_size=img_size,    patch_size=7,                   # stage1, 大卷积核7*7
    stride=4,                          # stage1, 4倍下采样
    in_chans=in_chans,    embed_dim=embed_dims[0])self.patch_embed2 = OverlapPatchEmbed(    img_size=img_size // 4,    patch_size=3,    stride=2,                          # stage2, 2倍下采样
    in_chans=embed_dims[0],    embed_dim=embed_dims[1])self.patch_embed3 = OverlapPatchEmbed(    img_size=img_size // 8,    patch_size=3,    stride=2,                          # stage3, 2倍下采样
    in_chans=embed_dims[1],    embed_dim=embed_dims[2])self.patch_embed4 = OverlapPatchEmbed(    img_size=img_size // 16,    patch_size=3,    stride=2,                          # stage4, 2倍下采样
    in_chans=embed_dims[2],    embed_dim=embed_dims[3])

class OverlapPatchEmbed(nn.Layer):
    def __init__(self,                 img_size=224,
                 patch_size=7,          # 卷积核大小
                 stride=4,                 # 下采样倍数
                 in_chans=3,            # 输入通道数
                 embed_dim=768):  # 输出通道数
        super().__init__()        img_size = to_2tuple(img_size)        patch_size = to_2tuple(patch_size)

        self.img_size = img_size
        self.patch_size = patch_size
        self.H, self.W = img_size[0] // patch_size[0], img_size[            1] // patch_size[1]
        self.num_patches = self.H * self.W        # 定义投影变换所用的卷积
        self.proj = nn.Conv2D(
            in_chans,
            embed_dim,            kernel_size=patch_size,
            stride=stride,
            padding=(patch_size[0] // 2, patch_size[1] // 2))        # 定义layer norm层
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):        x = self.proj(x)    # 通过卷积进行特征重投影，实现下采样、通道变换
        x_shape = paddle.shape(x)
        H, W = x_shape[2], x_shape[3]        x = x.flatten(2).transpose([0, 2, 1])  # 将H*W维度压缩成1个维度
        x = self.norm(x)          # 标准化

        return x, H, W

class Attention(nn.Layer):
    def __init__(self,                 dim,                 num_heads=8,                 qkv_bias=False,                 qk_scale=None,                 attn_drop=0.,                 proj_drop=0.,                 sr_ratio=1):
        super().__init__()
        assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."

        self.dim = dim
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5
        self.dim = dim
        
        # 定义q映射
        self.q = nn.Linear(dim, dim, bias_attr=qkv_bias)
        # 定义kv映射
        self.kv = nn.Linear(dim, dim * 2, bias_attr=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        # 定义输入特征的残差映射
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.sr_ratio = sr_ratio
        if sr_ratio > 1:
            self.sr = nn.Conv2D(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)     
            self.norm = nn.LayerNorm(dim)

    def forward(self, x, H, W):
        x_shape = paddle.shape(x)        B, N = x_shape[0], x_shape[1]        C = self.dim
        
        # 输入特征通过映射得到q
        q = self.q(x).reshape([B, N, self.num_heads,C // self.num_heads]).transpose([0, 2, 1, 3])
        
        # 输入特征通过映射得到k v
        if self.sr_ratio > 1:
            x_ = x.transpose([0, 2, 1]).reshape([B, C, H, W])
            x_ = self.sr(x_).reshape([B, C, -1]).transpose([0, 2, 1])         # 下采样
            x_ = self.norm(x_)
            kv = self.kv(x_).reshape([B, -1, 2, self.num_heads,C // self.num_heads]).transpose([2, 0, 3, 1, 4])
        else:
            kv = self.kv(x).reshape([B, -1, 2, self.num_heads,C // self.num_heads]).transpose([2, 0, 3, 1, 4])
        k, v = kv[0], kv[1]
        
        # att计算，q*k/sqrt(d)
        attn = (q @ k.transpose([0, 1, 3, 2])) * self.scale
        attn = F.softmax(attn, axis=-1)
        attn = self.attn_drop(attn)
        
        # att权重与x融合
        x = (attn @ v).transpose([0, 2, 1, 3]).reshape([B, N, C])
        # 关注后处理
        x = self.proj(x)
        x = self.proj_drop(x)

        return x

class Mlp(nn.Layer):
    def __init__(self,
                 in_features,
                 hidden_features=None,
                 out_features=None,
                 act_layer=nn.GELU,
                 drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.dwconv = DWConv(hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)    def forward(self, x, H, W):
        x = self.fc1(x)                    # 线性变换/MLP
        x = self.dwconv(x, H, W)  # 卷积/Conv3*3
        x = self.act(x)                   # GELU
        x = self.drop(x)
        x = self.fc2(x)                   # 线性变换/MLP
        x = self.drop(x)        return x

class SegFormer(nn.Layer):

    def __init__(self,
                 num_classes,
                 backbone,
                 embedding_dim,                 align_corners=False,
                 pretrained=None):
        super(SegFormer, self).__init__()

        self.pretrained = pretrained
        self.align_corners = align_corners
        self.backbone = backbone
        self.num_classes = num_classes
        c1_in_channels, c2_in_channels, c3_in_channels, c4_in_channels = self.backbone.feat_channels

        self.linear_c4 = MLP(input_dim=c4_in_channels, embed_dim=embedding_dim)
        self.linear_c3 = MLP(input_dim=c3_in_channels, embed_dim=embedding_dim)
        self.linear_c2 = MLP(input_dim=c2_in_channels, embed_dim=embedding_dim)
        self.linear_c1 = MLP(input_dim=c1_in_channels, embed_dim=embedding_dim)

        self.dropout = nn.Dropout2D(0.1)
        self.linear_fuse = layers.ConvBNReLU(            in_channels=embedding_dim * 4,            out_channels=embedding_dim,
            kernel_size=1,
            bias_attr=False)

        self.linear_pred = nn.Conv2D(
            embedding_dim, self.num_classes, kernel_size=1)

    def forward(self, x):        feats = self.backbone(x)
        c1, c2, c3, c4 = feats        ############## MLP decoder on C1-C4 ###########
        c1_shape = paddle.shape(c1)        c2_shape = paddle.shape(c2)        c3_shape = paddle.shape(c3)        c4_shape = paddle.shape(c4)        
        # 统一stage4的维度、尺寸
        _c4 = self.linear_c4(c4).transpose([0, 2, 1]).reshape([0, 0, c4_shape[2], c4_shape[3]])        _c4 = F.interpolate(
            _c4,            size=c1_shape[2:],
            mode='bilinear',
            align_corners=self.align_corners)
        
        # 统一stage3的维度、尺寸
        _c3 = self.linear_c3(c3).transpose([0, 2, 1]).reshape([0, 0, c3_shape[2], c3_shape[3]])        _c3 = F.interpolate(
            _c3,            size=c1_shape[2:],
            mode='bilinear',
            align_corners=self.align_corners)
        
        # 统一stage2的维度、尺寸
        _c2 = self.linear_c2(c2).transpose([0, 2, 1]).reshape([0, 0, c2_shape[2], c2_shape[3]])        _c2 = F.interpolate(
            _c2,            size=c1_shape[2:],
            mode='bilinear',
            align_corners=self.align_corners)
        
        # 统一stage1维度、尺寸
        _c1 = self.linear_c1(c1).transpose([0, 2, 1]).reshape(
            [0, 0, c1_shape[2], c1_shape[3]])        
        # 特征拼接与通道压缩
        _c = self.linear_fuse(paddle.concat([_c4, _c3, _c2, _c1], axis=1))
        logit = self.dropout(_c)        
        #分类预测
        logit = self.linear_pred(logit)
        return [
            F.interpolate(
                logit,                size=paddle.shape(x)[2:],
                mode='bilinear',
                align_corners=self.align_corners)
        ]

# pip升级!pip install --user --upgrade pip -i https://pypi.tuna.tsinghua.edu.cn/simple# 下载仓库，并切换到2.4版本%cd /home/aistudio/
!git clone https://gitee.com/paddlepaddle/PaddleSeg.git #该行仅在初次运行项目时运行即可，后续不需要运行改行命令%cd /home/aistudio/PaddleSeg
!git checkout -b release/2.4 origin/release/2.4# 安装依赖!pip install -r requirements.txt

# 耗时约35秒!unzip -oq /home/aistudio/data/data77571/train_and_label.zip -d /home/aistudio/data/src/
!unzip -oq /home/aistudio/data/data77571/img_test.zip -d /home/aistudio/data/src/

# 生产数据集划分txt# 演示时使用比例0.98：0.02!python /home/aistudio/work/segmentation/data_split.py \        0.98 0.02 0 \
        /home/aistudio/data/src/img_train \
        /home/aistudio/data/src/lab_train# # 实践时使用比例0.2：0.2# !python /home/aistudio/work/segmentation/data_split.py \#         0.8 0.2 0 \#         /home/aistudio/data/src/img_train \#         /home/aistudio/data/src/lab_train

# 修改文件!cp /home/aistudio/work/segmentation/segformerb3.yml /home/aistudio/PaddleSeg/segformerb3.yml
!cp /home/aistudio/work/segmentation/utils.py /home/aistudio/PaddleSeg/paddleseg/utils/utils.py                         # 加载tif数据与模型参数!cp /home/aistudio/work/segmentation/predict.py /home/aistudio/PaddleSeg/paddleseg/core/predict.py                      # 预测类别结果保存!cp /home/aistudio/work/segmentation/transformer_utils.py /home/aistudio/PaddleSeg/paddleseg/models/backbones/transformer_utils.py # 修复数据类型bug

# 演示时使用的训练超参数，约5分钟!python /home/aistudio/PaddleSeg/train.py \
    --config  /home/aistudio/PaddleSeg/segformerb3.yml \
    --save_dir /home/aistudio/data/output_seg \
    --do_eval \
    --use_vdl \
    --batch_size 32 \
    --iters 100 \
    --save_interval 50 \
    --log_iters 10 \
    --fp16 

# # 实践时使用的训练超参数，约20+小时# !python /home/aistudio/PaddleSeg/train.py \#     --config  /home/aistudio/PaddleSeg/segformerb3.yml \#     --save_dir /home/aistudio/data/output_seg \#     --do_eval \#     --use_vdl \#     --batch_size 32 \#     --iters 100000 \#     --save_interval 2100 \#     --log_iters 100 \#     --fp16

# 将训练参数转移到best_model/seg下!mkdir /home/aistudio/best_model
!mkdir /home/aistudio/best_model/seg
!cp /home/aistudio/data/output_seg/best_model/model.pdparams /home/aistudio/best_model/seg/model.pdparams

# 结果保存在/home/aistudio/data/infer_seg下!python /home/aistudio/PaddleSeg/predict.py \
       --config /home/aistudio/PaddleSeg/segformerb3.yml \
       --model_path /home/aistudio/best_model/seg/model.pdparams \
       --image_path /home/aistudio/data/src/img_testA \
       --save_dir /home/aistudio/data/infer_seg