本文为稀土技能社区首发签约文章，14天内制止转载，14天后未获授权制止转载，侵权必究！

作者简介：秃头小苏，致力于用最浅显的言语描绘问题

往期回顾：CV攻城狮入门VIT(vision transformer)之旅——近年超火的Transformer你再不了解就晚了！ 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的教程，这个模型才是目前真正霸榜的存在，敬请期待吧！！！

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