vit 代码学习

写在前面：对霹雳啪啦WZ大佬的代码进行学习，做一个日记，供复习用
更多可以访问大佬的链接：
github链接
博文地址

vit 代码学习

---

个人觉得 VIT模型的学习有两个难点，对于新手来说，一是attention模块，二就是图像块的切分。其余的其实和传统CNN模型差不多，搭好积木就可。attention模块可以看看transformer中的attention解释。
先放VIT的模型图：

在学习VIT 之前，首先要对transformer的self-attention机制有一个大致的了解

传统的transformer 模型：

self-attention 计算：

计算公式如下：

我们首先要理解第一个公式的意义：
在transformer的计算中：
在计算attention的时候，我们要计算q,k,v矩阵，但是q,k,v矩阵有什么意义呢，实际上我们可以考虑 X*transpose(X),其得到的值就是X中的每一行对它的第一行、第二行、直到最后一行的内积和（因为转置，Transpose(X)的列为X的行，故可以看作X的每一行和他其余行的关系重要程度，随后进行softmax操作，得到权重，再与原矩阵相乘就是加权后的结果了，即是第一个公式的意义。具体可参考：self-attention

vit中attention计算代码如下，假设输入图像尺寸为[8, 3,224,224]的图像，即一次输入8张图片，RGB 3通道图像，h,w为224大小，经过vit前面的patch 切分处理后，变为大小维度为[8,197,768] 大小，patch切分先按下不表。然后经过Drop-Out层，经过Transformer Encoder 进行注意力计算。程序如下：
还是按照上图的公式进行计算，首先线性层映射，最后一维度扩大3倍，然后矩阵维度变换，切分出q、k、v矩阵，最后按照公式进行计算即可。计算的意义可以参考上述链接。

class Attention(nn.Module):
    def __init__(self,
                 dim,   # 输入token的dim
                 num_heads=8,
                 qkv_bias=False,
                 qk_scale=None,
                 attn_drop_ratio=0.,
                 proj_drop_ratio=0.):
        super(Attention, self).__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads#多头注意力机制平分，768/12
        self.scale = qk_scale or head_dim ** -0.5 #qkv计算
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)#线性层，产生qkv
        self.attn_drop = nn.Dropout(attn_drop_ratio)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop_ratio)

    def forward(self, x):
        # [batch_size, num_patches + 1, total_embed_dim] 8*197*768
        B, N, C = x.shape

        # qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim] [8,197,3*768]
        # reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head] [8,197,3,12,768/12]
        # permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head] [3,8,12,197,768/12]
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple) [8,12,197,768/12]

        # transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1] [8,12,768/12,197]
        # @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1] 最后两维度矩阵乘法 [8,12,197,197]
        attn = (q @ k.transpose(-2, -1)) * self.scale # qkv计算公式
        attn = attn.softmax(dim=-1) #对行做softmax
        attn = self.attn_drop(attn)

        # @: multiply -> [batch_size, num_heads, num_patches + 1, embed_dim_per_head] [8,12,197,768/12]
        # transpose: -> [batch_size, num_patches + 1, num_heads, embed_dim_per_head]  [8,197,12,768/12]
        # reshape: -> [batch_size, num_patches + 1, total_embed_dim]              [8,197,768]
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

剩下的跟着模型图搭积木就可以了

图像块的划分

之前初步接触VIT、SwinT，一直不明白是怎么划分的，然后最后结果是卷积卷出来的。

class PatchEmbed(nn.Module):
    """
    2D Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None):
        super().__init__()
        img_size = (img_size, img_size)
        patch_size = (patch_size, patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])# 14*14 (224/16)
        self.num_patches = self.grid_size[0] * self.grid_size[1]#14*14

        self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)#输入：3*224*224 输出： （224-16）/16+1=14   768*14*14
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x):
        B, C, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."

        # flatten: [B, C, H, W] -> [B, C, HW]
        # transpose: [B, C, HW] -> [B, HW, C]
        x = self.proj(x).flatten(2).transpose(1, 2)
        x = self.norm(x)
        return x

重要代码其实就一句：

self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)

如：输入：3224224 输出： 7681414

其中768为输出的embed_dim，可以自己设置，长款经过卷积操作：（224-16）/16+1=14，14还有一个意义块的数量，卷积Stride 16还有一个意义是patch_size，图像块的尺寸，不过考虑到卷积也相当于一个特征融合操作，也就没有问题了。然后经过矩阵维度变化，就可以进行后续计算了。