当前位置：首页 > news >正文

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

news 2025/8/25 1:51:50

【图书推荐】《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,Blockclass 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_embedembedding = self.blocks(embedding)embedding = embedding[:,0]embedding = torch.nn.Dropout(0.1)(embedding)logits = self.logits_layer(embedding)return logitsif __name__ == '__main__':image = torch.randn(size=(2,3,84,84))VisionTransformer()(image)

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