ViT和SwinTransformer详解

ViT是Google brain发表于ICLR'21上的工作，开创性将transformer用在vision领域，且图像识别性能超CNN，至今引用3.8w+；原文：https://arxiv.org/pdf/2010.11929

SwinTransformer是微软亚洲研究院发表于ICCV'21上，获best paper，在多个视觉任务上获sota，打破CNN垄断vision backbone的现状，至今引用1.8w+；原文：https://openaccess.thecvf.com/content/ICCV2021/papers/Liu_Swin_Transformer_Hierarchical_Vision_Transformer_Using_Shifted_Windows_ICCV_2021_paper.pdf

建议读原文，这些文章优雅、简洁、深刻。

下面按照三部分进行，分别是Attention介绍、ViT详解、SwinTransformer详解。与常规文章讲解不同，我会多采用QA进行展开。

1. Attention介绍

这涉及到NeurIPS发表的“Attention is all you need”，这篇文章引用已经12w+，理解注意力机制是学习transformer的核心。

Q: general attention和self attention区别？

A: 相同点是均需要计算qkv，不同之处，self attention的input只有x，而general attention的input除了有x（映射得到kv），还有q(查询query)。

Self attention layer介绍

步骤：1. 输入x，通过映射矩阵Wq,Wk,Wv，得到qkv（D维）

2. q和k进行对齐操作，如：q0会分别与不同的k进行点乘操作，得e0（e矩阵第一列）

3. 注意力机制：softmax操作。如，从e0得a0（a矩阵第一列），为0~1之间的注意力权重

4. 输出：v和注意力权重a的加权和，如：y0为a0和所有v的注意力加权和

=》不同于CNN的局部特性，此处的自注意力很好地体现了全局特性。

仔细观察，可以发现self-attention layer具备permutation invariant的性质(置换不变)

现实中，不管是语言token还是vision patch token，位置不同，显然我们应该得到不同的内容向量y才是合理的。

因此，有必要加入位置编码，将位置信息考虑进来进行自注意力学习。

对于每个输入xj，给出位置编码pj。使用位置编码函数pos，pj=pos(j)，将位置j映射到D维向量（因为x是D维）。对于pos函数的选取此处不详细展开。

Multi-head self attention layer介绍

多头自注意力层就是transformer里核心模块。

Q：为什么要multi-head？

A：本质是为提取更好的特征，类似于CNN中卷积核也是多组，以得到多个特征谱。不同的是，CNN中卷积核小，计算量小，特征谱数量都是几十、几百。这儿的Multi-head不会很多，一般不超过10。

2. ViT详解

这篇文章的writing也可以当作范本，反复学习。

Q：标题两个keys，一个是an image is worth 16x16 words, 另一个是at scale，分别突出了什么？

A：前者突出将图像按照文字的处理方式，把一张图表示成了16x16 tokens。另一个关键点at scale，则与transformer的优势关联起来，也暗含了transformer要获的良好性能的前提。

Q：Transformer的天然优势是什么？

A：主要是excellent scalability，当模型和训练集增加时，并没有saturating performance。可以处理超大规模的训练数据。另一个是self-attention带来的computational efficiency，很多计算可以高度并行。

Q：CNN的天然优势是什么？

A：主要是inductive bias，在卷积的过程中，我们使用了translation equaivariance(平移不变性)、locality(局部性)来保留2D相邻结构。这些使得CNN在少量训练数据时候也能获得很好的性能。

Q：什么时候Transformer会比CNN更好？

A：通常，小训练数据集时候，convolutional inductive bias会很有用。当，数据集规模足够大的时候，最终large scaling training会比inductive bias表现好。这是合理的，因为Transformer学习中没有inductive bias，其特征时只能从大规模数据中学习。

Conclusion中提及的几个有前瞻性的点，现在均已经实现:)

1）self-supervised vs. supervised learning，之间的gap已经去掉；

2）scaling law，随着scaling提升，模型性能提升，现在已经是大模型发展遵循的发展规律；

3）transformer在segmentation、detection上的发展，现在已经横扫这些视觉任务。

# Transformer Encoder (depth x)

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Attention(dim, heads=heads, dim_head=dim_head,dropout=dropout),
                FeedForward(dim,mlp_dim,dropout=dropout)]))
    
    def forward(self, x):
        for att, ff in self.layers:
            x = attn(x)+x
            x = ff(x)+x
        return self.norm(x)

# Multi-Head Attention
# 与self-attention layer中的operation保持一致

class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropot=0.):
        super().__init__()
        inner_dim = dim_head * heads
        project_out = not(heads==1 and dim_head==dim)

        self.heads = heads
        self.scale = dim_head ** -0.5
        self.norm = nn.LayerNorm(dim)

        self.attend = nn.Softmax(dim=-1)    
        self.dropout = nn.Dropout(dropout)

        self.to_qkv = nn.Linear(dim, inner_dim*3, bias=False)
        
        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        x = self.norm(x)

        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q,k,v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv) # batch(b), sequence length(n), heads(h), dim(d) 
        dots = torch.matmul(q, k.transpose(-1,-2))*self.scale
        attn = self.attend(dots)
        attn = self.dropout(attn)
    
        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

# FeedForward (Transformer Encoder第二个部分)

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(nn.LayerNorm(dim),
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout))
        
        def forward(self, x):
            return self.net(x)

# ViT framework
# 包括输入图像的处理方式以及具体的任务

class ViT(nn.Module):
    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool='cls', channels=3, dim_head=64, dropout=0., emb_dropout=0.):
        super().__init__()
        image_height, image_width = pair(image_size)
        patch_height, patch_width = pair(patch_size)

        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
        
        num_pathes = (image_height//patch_height)*(image_width//patch_width) # sequence length
        patch_dim = channels*patch_height*patch_width
        assert pool in {'cls','mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)' #输入序列的第一个位置会添加一个特殊的标记，称为 [CLS] 标记

        self.to_patch_embedding = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
            nn.LayerNorm(patch_dim),
            nn.Linear(patch_dim, dim),
            nn.LayerNorm(dim),
        )
        
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches+1, dim))
        self.cls_token = nn.Parameter(torch.randn(1,1,dim))
        self.dropout = nn.Dropout(emb_dropout)

        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)

        self.pool = pool
        self.to_latent = nn.Identity()

        self.mlp_head = nn.Linear(dim, num_classes) # classification task

    def forward(self, img):
        x = self.to_patch_embedding(img)
        b,n,_ = x.shape

        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b=b)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.pos_embedding[:,:(n+1)]
        x = self.dropout(x)

        x = self.transformer(x)

        x = x.mean(dim=1) if self.pool=='mean' else x[:,0]

        x = self.to_latent(x)
        return self.mlp_head(x)

此处，pos_embedding是随机给的，transformer的输出后pool只能选cls或者mean中之一，然后进行MLP对任务的预测。

这里没有涉及到transformer decoder设计。

3. SwinTransformer详解

Q：ViT不好吗，SwinTransformer主要解决哪些关键问题？

A：如果图像分辨率变大，按照patch的size进行切分，这时候图像块的数量会增加，相应的计算复杂度quadratic增加，除此外切分的patch也相对较大（下采样倍数高），特征提取信息不准。不能很好处理高分辨率图。除此外，ViT这种固定的图像块切分方法对于不同大小的视觉实体而言不是很合理，当物体远小于或者大于patch时候，很难有效提取特征。不同物体的尺寸和比例差异很大，不像单词的长度相对固定。不能很好处理大小变化的视觉物体。

个人感觉SwinTransformer中窗口概念，类似与CNN中卷积核，窗口shift类似于CNN中stride，不同这个shift（向右、向下）更加灵活。不同在于，CNN针对局部图像感受野直接去求W，而SwinTransformer则是利用Self-attention更高级的方式去求局部图像的特征。

很多技巧都是用于减少（分窗口、窗口移动）运算量。

Q：SwinTransformer主要贡献？

A：第一，层级的特征谱方式使计算复杂度对于图像尺寸而言是linear而不是quadratic，可以处理高分辨率的图像。第二、shifted window很好解决了视觉物体大小变化的特点。

Q：SwinTransformer主要设计思想？

A：全局的注意力机制只在小范围内做，然后在不同层级上提特征（W-MSA，提出窗口的概念，窗口内进行多头注意力机制）。此外，利用shifted window将各个窗口之间的信息进行通信，完美达到捕获全局的上下文信息的优势（SW-MSA，此处就是滑动窗口的多头注意力机制）。这两部分就是Swin Transformer blocks的主要组成部分。