SwinTransformer 源码备忘

使用 "装饰器", 将 SwinTransformer 类注册到 BACKBONES 构建器.

@BACKBONES.register_module()
class SwinTransformer(nn.Module):
    ...

正向函数 forward:
x.shape : [1, 3, 800, 800], 通道 3, pixel 800x800

    self.patch_embed = PatchEmbed(patch_size=4, in_chans=3, embed_dim=96, norm_layer=True)
    ...
    def forward(self, x):
        x = self.patch_embed(x) # 将图像拆分为 patch

拆分函数 patch_embed
卷积 nn.Conv2d(3, 4, kernel_size=4, stride=4)
卷积, 拆分后: x.shape : [1, 96, 200, 200], 维度 96, patch 200x200

class PatchEmbed(nn.Module):
        ...
        self.proj = nn.Conv2d(3, 4, kernel_size=4, stride=4)
        ...
    def forward(self, x):
        ...
        x = self.proj(x)  # B C Wh Ww #

Norm - 标准化
返回: x.shape : [1, 96, 200, 200],

        # Norm
        Wh, Ww = x.size(2), x.size(3)
        x = x.flatten(2)
        x = x.transpose(1, 2)
        x = self.norm(x)
        x = x.transpose(1, 2)
        x = x.view(-1, self.embed_dim, Wh, Ww)

        return x

返回 SwinTransformer :: forward()
展平: x.shape : [1, 96, 40000]
变换: x.shape : [1, 40000, 96]

class SwinTransformer(nn.Module):
        ...
        x = x.flatten(2); 
        x = x.transpose(1, 2); 
        # x = self.pos_drop(x)

SwinTransformer 类定义, 包含 layers 列表

class SwinTransformer(nn.Module):
        ...
        # layers 初始化层
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = BasicLayer(...)
            self.layers.append(layer)

进入各层计算: x.shape : [1, 40000, 96]

class SwinTransformer(nn.Module):
    ...
    def forward(s