CV攻城狮入门VIT(vision transformer)之旅——VIT代码实战篇

作者简介：秃头小苏，致力于用最通俗的语言描述问题

专栏推荐：深度学习网络原理与实战

近期目标：写好专栏的每一篇文章

支持小苏：点赞、收藏⭐、留言

 

 

CV攻城狮入门VIT(vision transformer)之旅——VIT代码实战篇

写在前面

​​  在上一篇，我们已经介绍了VIT的原理，是不是发现还挺简单的呢！对VIT原理不清楚的请点击☞☞☞了解详细。那么这篇我将带大家一起来看看VIT的代码，主要为大家介绍VIT模型的搭建过程，也会简要的说说训练过程。

​​  这篇VIT的模型是用于物体分类的，我们选择的例子是花的五分类问题。关于花的分类，我之前也有详细的介绍，是用卷积神经网络实现的，不清楚可以点击下列链接了解详情：

基于pytorch搭建AlexNet神经网络用于花类识别

基于pytorch搭建VGGNet神经网络用于花类识别

基于pytorch搭建GoogleNet神经网络用于花类识别

基于pytorch搭建ResNet神经网络用于花类识别

​​  代码部分依旧参考的是B站霹雳吧啦Wz 的视频 ，强烈推荐大家观看喔，你一定会收获满满！！！如果你看视频中有什么不理解的，可以来这篇文章寻找寻找答案喔。

​  代码点击☞☞☞获取。

 
 

VIT模型构建

​  这部分我以VIT-Base模型为例为大家讲解，此模型的相关参数如下：

Model	Patch size	Layers	Hidden Size	MLP size	Heads	Params
VIT-Base	16*16	12	768	3072	12	86M

​​  在上代码之前，我们有必要了解整个VIT模型的结构。关于这点我在上一篇VIT原理详解篇已经为大家介绍过，但上篇模型结构上的一些细节，像Droupout层，Encoder结构等等都是没有体现的，这些只有阅读源码才知道。下面给出整个VIT-Base模型的详细结构，如下图所示：

​

图片来自于 霹雳吧啦Wz 的 博客

​​  我们的代码是完全按照上图结构搭建的，但在解读代码之前我觉得很有必要再向大家强调一件事——你看我上文推荐的视频或看我的代码解读都只起到一个辅助的作用，你很难说光靠看就能把这些理解透彻。我当时看视频的时候甚至很难完整的看完一遍，更多的还是靠自己一步一步的调试来看每个操作后维度的变换。

​​  我猜测可能有些同学还不是很清楚怎么在vit_model.py进行调试，其实很简单，只需要创建一个全1的tensor来模拟图片，将其当作输入输入网络即可，即可在vit_model.py文件末尾加上下列代码：

if __name__ == '__main__':
    input = torch.ones(1, 3, 224, 224)    # 1为batch_size   (3 224 224)即表示输入图片尺寸
    print(input.shape)
    model = vit_base_patch16_224_in21k()  #使用VIT_Base模型,在imageNet21k上进行预训练
    output = model(input)
    print(output.shape)

---

---

​​  那么下面我们就一步步的对代码进行解读，首先我们先对输入进行Patch_embedding操作，这部分我在理论详解篇有详细的介绍过，其就是采用一个卷积核大小为16*16，步长为16的卷积和一个展平操作实现的，相关代码如下：

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])
        self.num_patches = self.grid_size[0] * self.grid_size[1]

        self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)
        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

​​  其实我觉得我再怎么解释这个代码的效果都不会很好，你只要在这里打上一个断点，这个过程就一目了然了。所以这篇文章可能就更倾向于让大家熟悉一下整个模型搭建的过程，具体细节大家可自行调试！！！

---

​  这步结束后，你会发现现在x的维度为(1，196，768)。其中1为batch_size数目，我们之前将其设为1。

​  接着我们会将此时的x和Class token拼接，相关代码如下：

# 定义一个可学习的Class token
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))  # 第一个1为batch_size   embed_dim=768 
cls_token = self.cls_token.expand(x.shape[0], -1, -1)        # 保证cls_token的batch维度和x一致
if self.dist_token is None:
    x = torch.cat((cls_token, x), dim=1)  # [B, 197, 768]    self.dist_token为None,会执行这句
else:
    x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)

​  同样可以来看看拼接后的维度，如下图：

---

​  继续进行下一步——位置编码。位置编码是和上步得到的x进行相加的操作，相关代码如下:

 # 定义一个可学习的位置编码
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))   #这个维度为（1，197，768）
x = x + self.pos_embed

​​  经过位置编码输入的维度并不会发生变换，如下：

​​  位置编码过后，还会经过一个Dropout层，这并不会改变输入维度，相信大家对这个就很熟悉了，就不过多介绍了。

---

​  到这里，我们的输入维度为（1，197，768）。接下来就要被送入encoder模块了。首先做了一个Layer Normalization归一化操作，接着会送入Multi-Head Attention部分，然后进行Droppath操作并做一个残差链接。这部分的代码如下：

class Block(nn.Module):
    def __init__(self,
                 dim,
                 num_heads,
                 mlp_ratio=4.,
                 qkv_bias=False,
                 qk_scale=None,
                 drop_ratio=0.,
                 attn_drop_ratio=0.,
                 drop_path_ratio=0.,
                 act_layer=nn.GELU,
                 norm_layer=nn.LayerNorm):
        super(Block, self).__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                              attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)

    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x)))   #上文描述的在这喔
        x = x + self.drop_path(self.mlp(self.norm2(x)))    #这是encode结构的后半部分
        return x

​  相信你对Layer Normalization已经有相关了解了，不清楚的可以看我对Transfomer讲解的文章，里面有关于此部分的解释，这里不再重复叙述。但是你对Multi-Head Attention是如何实现的可能还存在诸多疑惑，此部代码如下：

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
        self.scale = qk_scale or head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        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]
        B, N, C = x.shape

        # qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim]
        # reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head]
        # permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]
        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)

        # transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1]
        # @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1]
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

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

​  光看确实难以发现其中的很多细节，那就尽情的调试吧！！！这部分也不会改变x的尺寸，如下：

​​  Multi-Head Attention后还有个Droppath层，其和Dropout类似，但说实话我也没了解过，就当成是一个固定的模块使用了。感兴趣的可以查阅资料。如果有很多人不了解或者我后期会经常用到这个函数的话，我也会出一期Dropout和Droppath区别的教程。这里就靠大家自己啦！！！

​  下一步同样是一个Layer Normalization层，接着是MLP Block，最后是一个Droppath加一个残差链接。这一部分还值得说的就是这个MLP Bolck了，但其实也非常简单，主要就是两个全连接层，相关代码如下：

class Mlp(nn.Module):
    """
    MLP as used in Vision Transformer, MLP-Mixer and related networks
    """
    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.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

​​  需要提醒大家的是上述代码的hidden_features其实就是一开始模型参数中MLP size，即3072。

---

​​  这样一个encoder Block就介绍完了，接着只需要重复这个Block 12次即可。这部分相关代码如下：

self.blocks = nn.Sequential(*[
            Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                  drop_ratio=drop_ratio, attn_drop_ratio=attn_drop_ratio, drop_path_ratio=dpr[i],
                  norm_layer=norm_layer, act_layer=act_layer)
            for i in range(depth)
        ])
        
        
x = self.blocks(x)

---

---

​  注意输入输出这个encoder Block前后，x的维度同样没有发生变化，仍为（1，197，768）。接着会进行Layer Normalization操作。然后要通过切片的方式提取出Class Token，代码如下：

if self.dist_token is None:
    return self.pre_logits(x[:, 0])    #self.dist_token=None  执行此句
 else:
    return x[:, 0], x[:, 1]

​  你会发现上述代码中会存在一个pre_logits()函数，这个函数其实就是一个全连接层加上一个Tanh激活函数，如下：

# Representation layer
if representation_size and not distilled:
    self.has_logits = True
    self.num_features = representation_size
    self.pre_logits = nn.Sequential(OrderedDict([
        ("fc", nn.Linear(embed_dim, representation_size)),
        ("act", nn.Tanh())
    ]))
else:
    self.has_logits = False
    self.pre_logits = nn.Identity()

​  可以发现，这部分不是总存在的。当representation_size=None时，此部分只是一个恒等映射，即什么都不做。关于representation_size何时取何值，我这里做一个简要的说明。当我们的预训练数据集是ImageNet时，representation_size=None，即此时什么都不做；当预训练数据集为ImageNet-21k时，representation_size是一个特定的值，至于是多少是不定的，这和是Base、Large或Huge模型有关，我们这里以Base模型为例，representation_size=768。

​​  经过pre_logits后，还有最后一个全连接层用于最终的分类。相关代码如下：

self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
x = self.head(x)

​​  到这里，VIT模型的搭建就全部介绍完啦，看到这里的话，为自己鼓个掌吧

 
 

VIT 训练脚本

​​  VIT训练部分和之前我用神经网络搭建的花类识别训练脚本基本是一样的，不清楚的可以先去看看之前的文章。这里我给大家讲讲怎么进行训练。其实你需要修改的地方只有两处，第一是数据集的路径，在代码中设置默认路径如下：

 parser.add_argument('--data-path', type=str,
                        default="/data/flower_photos")

​​  我们只需要将"/data/flower_photos"修改成我们对应的数据集路径即可。需要注意的是这里路径要指定到flower_photos文件夹，否则检测不到图片，这里和之前讲的还是有点差别的。

​  还有一处你需要修改的地方为预训练权重的位置，代码中默认路径如下：

# 预训练权重路径，如果不想载入就设置为空字符
parser.add_argument('--weights', type=str, default='./vit_base_patch16_224_in21k.pth',
                    help='initial weights path')

​  我们需要将'./vit_base_patch16_224_in21k.pth'换成自己下载预训练权重的地址。需要注意的时这里的预训练权重需要和你创建模型时选择的模型是一样的，即你选择了VIT_Base模型并在ImageNet21k上做预训练，你就要使用./vit_base_patch16_224_in21k.pth的预训练权重。

​​  最后我们训练的权重会保存在当前文件夹下的weights文件夹下，没有这个文件夹会创建一个新的，相关代码如下：

torch.save(model.state_dict(), "./weights/model-{}.pth".format(epoch))

 
 

VIT分类任务实验结果

​ 这里我们来看看花的五分类训练结果：

不使用预训练模型训练10轮：

不使用预训练权重训练50轮：

使用预训练权重训练10轮：

​  通过上面的三个实验你可以发现，VIT模型不使用预训练权重进行训练的话效果是非常差的，我们用ResNet网络不使用预训练权重训练50轮大概能达到0.79左右的准确率，而ViT只能达到0.561；但是使用了预训练模型的ResNet达到了0.915，而VIT高达0.971，效果是非常不错的。所以VIT是非常依赖预训练的，且预训练数据集越大，效果往往越好。

​  最后我们来看看预测部分，下图为检测郁金香的概率：

 
 

小结

​​  到这里，VIT代码实战篇就介绍完了。同时CV攻城狮入门VIT(vision transformer)之旅的三篇文章到这里也就告一个段落了，希望大家能够有所收获吧！！！

​​  这里预告一下，后期我打算出Swin Transformer的教程，这个模型才是目前真正霸榜的存在，敬请期待吧！！！

 
 
如若文章对你有所帮助，那就