Swin-Unet最强分割网络

Swin-Unet是基于Swin Transformer为基础(可参考Swin Transformer介绍 )，结合了U-Net网络的特点(可参考Tensorflow深度学习算法整理(三) 中的U-Net)组合而成的新的分割网络

它与Swin Transformer不同的地方在于，在编码器(Encoder)这边虽然跟Swin Transformer一样的4个Stage，但Swin Transformer Block的数量为[2,2,2,1]，而不是Swin Transformer的[2,2,6,2]。而在解码器(Decoder)这边,由于是升采样，使用的不再是Patch Embedding和Patch Merging，而使用的是Patch Expanding，它是Patch Merging的逆过程。

我们来看一下Patch Expanding的代码实现

from einops import rearrange

class PatchExpand(nn.Module):
    """
    块状扩充，尺寸翻倍，通道数减半
    """
    def __init__(self, input_resolution, dim, dim_scale=2, norm_layer=nn.LayerNorm):
        """
        Args:
            input_resolution: 解码过程的feature map的宽高
            dim: frature map通道数
            dim_scale: 通道数扩充的倍数
            norm_layer: 通道方向归一化
        """
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        # 通过全连接层来扩大通道数
        self.expand = nn.Linear(dim, 2 * dim, bias=False) if dim_scale == 2 else nn.Identity()
        self.norm = norm_layer(dim // dim_scale)

    def forward(self, x):
        """
        x: B, H*W, C
        """
        H, W = self.input_resolution
        # 先把通道数翻倍
        x = self.expand(x)
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"

        x = x.view(B, H, W, C)
        # 将各个通道分开，再将所有通道拼成一个feature map
        # 增大了feature map的尺寸
        x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=2, p2=2, c=C // 4)
        # 通道翻倍后再除以4，实际相当于通道数减半
        x = x.view(B, -1, C // 4)
        x = self.norm(x)

        return x

在编码器这边基本上跟Swin Transformer是一样的，我们重点来看解码器这边。它是使用BasicLayer_up类来对SwinTransformerBlock和Patch Expanding来进行搭配的。

class BasicLayer_up(nn.Module):
    """ A basic Swin Transformer layer for one stage.
    一个BasicLayer_up包含偶数个SwinTransformerBlock和一个upsamele层（即Patch Expanding层）
    """

    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., norm_layer=nn.LayerNorm, upsample=None, use_checkpoint=False):
        """
        Args:
            dim: feature map通道数
            input_resolution: feature map的宽高
            depth: 各个Stage中，Swin Transformer Block的数量
            num_heads: 多头注意力各个Stage中的头数
            window_size: 窗口自注意力机制的窗口中的patch数
            mlp_ratio: 层感知机模块中第一个全连接层输出的通道倍数
            qkv_bias: 如果是True的话，对自注意力公式中的Q、K、V增加一个可学习的偏置
            qk_scale: 窗口自注意力公式常数
            drop: dropout rate，默认为0
            attn_drop: 用于自注意力机制中的dropout rate，默认为0
            drop_path: 在Swin Transformer Block中，有一定概率丢弃整个直连分支，包括
                       LN、W-MSA或者SW-MSA，只保留直连的连接，是一种网络深度的随机性，默认为0
            norm_layer: 通道方向归一化
            upsample: 使用Patch Expanding来升采样
            use_checkpoint: 是否使用Pytorch中间数据保存机制
        """

        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint

        # build SwinTransformerBlock
        self.blocks = nn.ModuleList([
            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
                                 num_heads=num_heads, window_size=window_size,
                                 # 用于区分是使用W-MSA还是SW-MSA，0为W-MSA，1为SW-MSA
                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
                                 mlp_ratio=mlp_ratio,
                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
                                 drop=drop, attn_drop=attn_drop,
                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
                                 norm_layer=norm_layer)
            for i in range(depth)])

        # patch merging layer
        # 当stage=4的时候为None
        if upsample is not None:
            self.upsample = PatchExpand(input_resolution, dim=dim, dim_scale=2, norm_layer=norm_layer)
        else:
            self.upsample = None

    def forward(self, x):
        # 通过每一个SwinTransformerBlock
        for blk in self.blocks:
            if self.use_checkpoint:
                x = checkpoint.checkpoint(blk, x)
            else:
                x = blk(x)
        # 进行块状扩充(PatchExpanding)上采样
        if self.upsample is not None:
            x = self.upsample(x)
        return x

SwinTransformerBlock跟SwinTransformer中的代码也是一样的，这里就不重复了。

然后还有一个从编码器到解码器之间的跳连。这里需要看一下Swin-Unet的主类代码

class SwinTransformerSys(nn.Module):
    """ Swin-UNet网络模型
    """

    def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
                 embed_dim=96, depths=[2, 2, 2, 2], depths_decoder=[1, 2, 2, 2], num_heads=[3, 6, 12, 24],
                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
                 use_checkpoint=False, final_upsample="expand_first", **kwargs):
        """
        Args:
            img_size: 原始图像尺寸
            patch_size: 一个patch中的像素点数
            in_chans: 进入网络的图片通道数
            num_classes: 分类数量
            embed_dim: feature map通道数
            depths: 编码器各个Stage中，Swin Transformer Block的数量
            depths_decoder: 解码器各个Stage中，Swin Transformer Block的数量
            num_heads: 多头注意力各个Stage中的头数
            window_size: 窗口自注意力机制的窗口中的patch数
            mlp_ratio: 多层感知机模块中第一个全连接层输出的通道倍数
            qkv_bias: 如果是True的话，对自注意力公式中的Q、K、V增加一个可学习的偏置
            qk_scale: 自注意力公式中的常量
            drop_rate: dropout rate，默认为0
            attn_drop_rate: 用于自注意力机制中的dropout rate，默认为0
            drop_path_rate: 在Swin Transformer Block中，有一定概率丢弃整个直连分支，包括
                            LN、W-MSA或者SW-MSA，只保留直连的连接，是一种网络深度的随机性，默认为0.1
            norm_layer: 通道方向归一化
            ape: 是否进行绝对位置嵌入，默认False
            patch_norm: 如果是True的话，在patch embedding之后加上归一化
            use_checkpoint: 是否使用Pytorch中间数据保存机制
            final_upsample: 解码器stage4后的Patch Expanding
            **kwargs:
        """
        super().__init__()

        print("SwinTransformerSys expand initial----depths:{};depths_decoder:{};drop_path_rate:{};num_classes:{}".format(depths,
              depths_decoder, drop_path_rate, num_classes))

        self.num_classes = num_classes
        # stage的数量
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.ape = ape
        self.patch_norm = patch_norm
        # 编码器stage4输出特征的通道数(Swin-Tiny:768)
        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
        # 解码器stage4输出特征的通道数(192)
        self.num_features_up = int(embed_dim * 2)
        self.mlp_ratio = mlp_ratio
        self.final_upsample = final_upsample

        # 把图像分割成不重叠的patch
        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None)
        num_patches = self.patch_embed.num_patches
        # 获取feature map的高宽
        patches_resolution = self.patch_embed.patches_resolution
        self.patches_resolution = patches_resolution

        # 绝对位置嵌入
        if self.ape:
            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
            trunc_normal_(self.absolute_pos_embed, std=.02)

        self.pos_drop = nn.Dropout(p=drop_rate)

        # 不同的stage，舍弃整个直连分支的概率不同，从小到大，最小为0，最大为0.1
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule

        # 创建编码器layers
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):  # layer相当于stage
            layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer),
                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
                                                 patches_resolution[1] // (2 ** i_layer)),
                               depth=depths[i_layer],
                               num_heads=num_heads[i_layer],
                               window_size=window_size,
                               mlp_ratio=self.mlp_ratio,
                               qkv_bias=qkv_bias, qk_scale=qk_scale,
                               drop=drop_rate, attn_drop=attn_drop_rate,
                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
                               norm_layer=norm_layer,
                               # 只有前3个stage有patchmerging，最后一个没有
                               downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
                               use_checkpoint=use_checkpoint)
            self.layers.append(layer)
        
        # 创建解码器layers
        self.layers_up = nn.ModuleList()
        self.concat_back_dim = nn.ModuleList()
        for i_layer in range(self.num_layers):  # layer相当于stage
            # 每一个stage结束后，通道数减半的全连接层
            concat_linear = nn.Linear(2 * int(embed_dim * 2**(self.num_layers - 1 - i_layer)),
                                      int(embed_dim * 2**(self.num_layers - 1 - i_layer))) if i_layer > 0 else nn.Identity()
            if i_layer == 0:  # 第一个stage只进行上采样
                layer_up = PatchExpand(input_resolution=(patches_resolution[0] // (2 ** (self.num_layers - 1 - i_layer)),
                                       patches_resolution[1] // (2 ** (self.num_layers-1-i_layer))), dim=int(embed_dim * 2 ** (self.num_layers-1-i_layer)), dim_scale=2, norm_layer=norm_layer)
            else:
                layer_up = BasicLayer_up(dim=int(embed_dim * 2 ** (self.num_layers-1-i_layer)),
                                         input_resolution=(patches_resolution[0] // (2 ** (self.num_layers-1-i_layer)),
                                                           patches_resolution[1] // (2 ** (self.num_layers-1-i_layer))),
                                         depth=depths[(self.num_layers-1-i_layer)],
                                         num_heads=num_heads[(self.num_layers-1-i_layer)],
                                         window_size=window_size,
                                         mlp_ratio=self.mlp_ratio,
                                         qkv_bias=qkv_bias, qk_scale=qk_scale,
                                         drop=drop_rate, attn_drop=attn_drop_rate,
                                         drop_path=dpr[sum(depths[:(self.num_layers-1-i_layer)]):sum(depths[:(self.num_layers - 1 - i_layer) + 1])],
                                         norm_layer=norm_layer,
                                         # 只有前3个stage有PatchExpand，最后一个没有
                                         upsample=PatchExpand if (i_layer < self.num_layers - 1) else None,
                                         use_checkpoint=use_checkpoint)
            self.layers_up.append(layer_up)
            self.concat_back_dim.append(concat_linear)

        self.norm = norm_layer(self.num_features)
        self.norm_up = norm_layer(self.embed_dim)
        # 解码器最后一个stage进行FinalPatchExpand处理
        if self.final_upsample == "expand_first":
            print("---final upsample expand_first---")
            self.up = FinalPatchExpand_X4(input_resolution=(img_size // patch_size, img_size // patch_size), dim_scale=4, dim=embed_dim)
            self.output = nn.Conv2d(in_channels=embed_dim, out_channels=self.num_classes, kernel_size=1, bias=False)

        self.apply(self._init_weights)

这里有一个FinalPatchExpand_X4的方法，我们来看一下它的实现

class FinalPatchExpand_X4(nn.Module):
    """
    stage4之后的PatchExpand
    尺寸翻倍，通道数不变
    """
    def __init__(self, input_resolution, dim, dim_scale=4, norm_layer=nn.LayerNorm):
        """
        Args:
            input_resolution: feature map的宽高
            dim: feature map通道数
            dim_scale: 通道数扩充的倍数
            norm_layer: 通道方向归一化
        """
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.dim_scale = dim_scale
        # 通过全连接层来扩大通道数
        self.expand = nn.Linear(dim, 16 * dim, bias=False)
        self.output_dim = dim 
        self.norm = norm_layer(self.output_dim)

    def forward(self, x):
        """
        x: B, H*W, C
        """
        H, W = self.input_resolution
        # 先把通道数翻倍
        x = self.expand(x)
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"

        x = x.view(B, H, W, C)
        # 将各个通道分开，再将所有通道拼成一个feature map
        # 增大了feature map的尺寸
        x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=self.dim_scale, p2=self.dim_scale, c=C//(self.dim_scale**2))
        # 把扩大的通道数转成原来的通道数
        x = x.view(B, -1, self.output_dim)
        x = self.norm(x)

        return x

回到SwinTransformerSys代码中

def _init_weights(self, m):
    """
    对全连接层或者通道归一化进行权重以及偏置的初始化
    """
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.LayerNorm):
        nn.init.constant_(m.bias, 0)
        nn.init.constant_(m.weight, 1.0)

@torch.jit.ignore
def no_weight_decay(self):
    return {'absolute_pos_embed'}

@torch.jit.ignore
def no_weight_decay_keywords(self):
    return {'relative_position_bias_table'}

#Encoder and Bottleneck
def forward_features(self, x):
    """
    编码器过程
    """
    # 图像分割
    x = self.patch_embed(x)
    # 绝对位置嵌入
    if self.ape:
        x = x + self.absolute_pos_embed
    x = self.pos_drop(x)
    # 跳连点
    x_downsample = []
    # 通过各个编码过程的stage
    for layer in self.layers:
        x_downsample.append(x)
        x = layer(x)

    x = self.norm(x)  # B L C

    return x, x_downsample

#Dencoder and Skip connection
def forward_up_features(self, x, x_downsample):
    """
    解码器过程，包含了跳连拼接
    """
    # 通过各个解码过程的stage
    for inx, layer_up in enumerate(self.layers_up):
        if inx == 0:
            x = layer_up(x)
        else:
            # 拼接编码器的跳连部分再进入Swin Transformer Block
            x = torch.cat([x, x_downsample[3-inx]], -1)
            x = self.concat_back_dim[inx](x)
            x = layer_up(x)

    x = self.norm_up(x)  # B L C

    return x

def up_x4(self, x):
    """
    完成解码器的最后一个stage后进入
    """
    H, W = self.patches_resolution
    B, L, C = x.shape
    assert L == H * W, "input features has wrong size"

    if self.final_upsample == "expand_first":
        x = self.up(x)
        x = x.view(B, 4 * H, 4 * W, -1)
        x = x.permute(0, 3, 1, 2) #B,C,H,W
        x = self.output(x)
        
    return x

def forward(self, x):
    """
    前向运算
    """
    x, x_downsample = self.forward_features(x)
    x = self.forward_up_features(x, x_downsample)
    x = self.up_x4(x)

    return x

def flops(self):
    flops = 0
    flops += self.patch_embed.flops()
    for i, layer in enumerate(self.layers):
        flops += layer.flops()
    flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)
    flops += self.num_features * self.num_classes
    return flops