基于Vision Transformer的mini_ImageNet图片分类实战

【图书推荐】《PyTorch深度学习与计算机视觉实践》-CSDN博客

PyTorch计算机视觉之Vision Transformer 整体结构-CSDN博客

mini_ImageNet数据集简介与下载

mini_ImageNet数据集节选自ImageNet数据集。ImageNet是一个非常有名的大型视觉数据集，它的建立旨在促进视觉识别研究。ImageNet为超过1400万幅图像进行了注释，而且给至少100万幅图像提供了边框。同时，ImageNet包含2万多个类别，比如“气球”“轮胎”和“狗”等类别，ImageNet的每个类别均不少于500幅图像。

训练这么多图像需要消耗大量的资源，为了节约资源，后续的研究者在全ImageNet的基础上提取出了mini_ImageNet数据集。Mini_ImageNet包含100类共60000幅彩色图片，其中每类有600个样本，每幅图片的规格为84×84。通常而言，这个数据集的训练集和测试集的类别划分为80:20。相比于CIFAR-10数据集，mini_ImageNet数据集更加复杂，但更适合进行原型设计和实验研究。

mini_ImageNet的下载也很容易，读者可以使用提供的库包完成对应的下载操作，安装命令如下：

pip install MLclf

Vision Transformer模型设计

下面就是对训练过程的Vision Transformer进行模型设计，在11.1.4节完成的Vision Transformer模型的设计，针对的是224维度大小的图片，而此时使用的是mini版本的ImageNet，因此在维度上会有所变换。本例Vision Transformer模型的完整代码如下：

import torch
from vit import PatchEmbed,Block

class VisionTransformer(torch.nn.Module):
    def __init__(self,num_patches = 1,image_size = 84,patch_size = 14,embed_dim = 588,num_heads = 6,
                 qkv_bias = True,depth = 3,num_class = 64):
        super().__init__()

        #初始化PatchEmbed层
        self.patch_embed  = PatchEmbed(img_size = image_size,patch_size=patch_size,embed_dim=embed_dim)
        #增加一个作为标志物的参数
        self.cls_token = torch.nn.Parameter(torch.zeros(1, 1, embed_dim))

        #建立位置向量，计算embedding的长度
        self.num_tokens = (image_size * image_size) // (patch_size * patch_size)
        self.pos_embed = torch.nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))

        #这里在使用block模块时采用了指针的方式，注意*号
        self.blocks = torch.nn.Sequential(
            *[Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=4.0, qkv_bias=qkv_bias) for _ in range(depth)]
        )
        #最终的logits推断层
        self.logits_layer = torch.nn.Sequential(torch.nn.Linear(embed_dim, 512),torch.nn.GELU(),torch.nn.Linear(512, num_class))

    def forward(self,x):

        embedding = self.patch_embed(x)

        #添加标志物
        cls_token = self.cls_token.expand(x.shape[0], -1, -1)
        embedding = torch.cat((cls_token, embedding), dim=1)  #[B, 197, 768]
        embedding += self.pos_embed

        embedding = self.blocks(embedding)

        embedding = embedding[:,0]
        embedding = torch.nn.Dropout(0.1)(embedding)
        logits = self.logits_layer(embedding)
        return logits

if __name__ == '__main__':
    image = torch.randn(size=(2,3,84,84))
    VisionTransformer()(image)

《PyTorch深度学习与计算机视觉实践（人工智能技术丛书）》(王晓华)【摘要 书评 试读】- 京东图书 (jd.com)