devbox
weixin_40725706
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

【pytorch】Vision Transformer实现图像分类+可视化+训练数据保存_vision transformer 训练自己的数据

作者：weixin_40725706 | 2024-04-04 14:15:02

赞

踩

vision transformer 训练自己的数据

一、Vision Transformer介绍

Transformer的核心是 “自注意力” 机制。

论文地址：https://arxiv.org/pdf/2010.11929.pdf

自注意力（self-attention）相比 卷积神经网络 和 循环神经网络 同时具有并行计算和最短的最大路径⻓度这两个优势。因此，使用自注意力来设计深度架构是很有吸引力的。对比之前仍然依赖循环神经网络实现输入表示的自注意力模型 [Cheng et al., 2016,Lin et al., 2017b, Paulus et al., 2017]，transformer模型完全基于注意力机制，没有任何卷积层或循环神经网络层 [Vaswani et al., 2017]。尽管transformer最初是应用于在文本数据上的序列到序列学习，但现在已经推广到各种现代的深度学习中，例如语言、视觉、语音和强化学习领域。

17年发布时主要应用于不同语言之间翻译功能的实现。而在后来，有关研究发现Transformer应用于计算机视觉CV方面有着不输于卷积神经网络的强劲性能，一定程度上甚至比卷积神经网络更强。于是，初代Vision Transformer诞生了， 简称Vit。

Vision Transformer和Transformer区别是什么？用最最最简单的理解方式来看，Transformer的工作就是把一句话从一种语言翻译成另一种语言。主要是通过是将待翻译的一句话拆分为 多个单词 或者 多个模块，进行编码和解码训练，再评估那个单词对应的意思得分高就是相应的翻译结果。

而Vision Transformer则是将一个图片抽象地看做翻译中一个句子，通过图像分割将其拆分为多个模块，再进行编码和解码训练，评估中得分高的选项便是预测的结果。（纯属个人理解，如有错误，欢迎批评指正）

二、数据集

我的数据集为植物叶片病害的无标注数据集，共有三种类型。

  1. {
  2.     "0": "Huanglong_disease",
  3.     "1": "Magnesium_deficiency",
  4.     "2": "Normal"
  5. }

其中train : val : test  =  8 : 1 : 1，种类都是三种，只是数量不一样。

  1. train
  2. ├── Huanglong_disease
  3. │    ├── 000000.jpg
  4. │    ├── 000001.jpg
  5. │    ├── 000002.jpg
  6. │    ├── .............
  7. │    ├── 000607.jpg
  8. ├── Magnesium_deficiency
  9. └── Normal

大概长这样：

三、实战代码

1.vit_model.py

  1. """
  2. original code from rwightman:
  3. https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
  4. """
  5. from functools import partial
  6. from collections import OrderedDict
  7.  
  8. import torch
  9. import torch.nn as nn
  10.  
  11.  
  12. def drop_path(x, drop_prob: float = 0., training: bool = False):
  13.     """
  14.     Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
  15.     This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
  16.     the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
  17.     See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
  18.     changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
  19.     'survival rate' as the argument.
  20.     """
  21.     if drop_prob == 0. or not training:
  22.         return x
  23.     keep_prob = 1 - drop_prob
  24.     shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
  25.     random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
  26.     random_tensor.floor_()  # binarize
  27.     output = x.div(keep_prob) * random_tensor
  28.     return output
  29.  
  30.  
  31. class DropPath(nn.Module):
  32.     """
  33.     Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
  34.     """
  35.     def __init__(self, drop_prob=None):
  36.         super(DropPath, self).__init__()
  37.         self.drop_prob = drop_prob
  38.  
  39.     def forward(self, x):
  40.         return drop_path(x, self.drop_prob, self.training)
  41.  
  42.  
  43. class PatchEmbed(nn.Module):
  44.     """
  45.     2D Image to Patch Embedding
  46.     """
  47.     def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None):
  48.         super().__init__()
  49.         img_size = (img_size, img_size)
  50.         patch_size = (patch_size, patch_size)
  51.         self.img_size = img_size
  52.         self.patch_size = patch_size
  53.         self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
  54.         self.num_patches = self.grid_size[0] * self.grid_size[1]
  55.  
  56.         self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)
  57.         self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
  58.  
  59.     def forward(self, x):
  60.         B, C, H, W = x.shape
  61.         assert H == self.img_size[0] and W == self.img_size[1], \
  62.             f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
  63.  
  64.         # flatten: [B, C, H, W] -> [B, C, HW]
  65.         # transpose: [B, C, HW] -> [B, HW, C]
  66.         x = self.proj(x).flatten(2).transpose(1, 2)
  67.         x = self.norm(x)
  68.         return x
  69.  
  70.  
  71. class Attention(nn.Module):
  72.     def __init__(self,
  73.                  dim,   # 输入token的dim
  74.                  num_heads=8,
  75.                  qkv_bias=False,
  76.                  qk_scale=None,
  77.                  attn_drop_ratio=0.,
  78.                  proj_drop_ratio=0.):
  79.         super(Attention, self).__init__()
  80.         self.num_heads = num_heads
  81.         head_dim = dim // num_heads
  82.         self.scale = qk_scale or head_dim ** -0.5
  83.         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
  84.         self.attn_drop = nn.Dropout(attn_drop_ratio)
  85.         self.proj = nn.Linear(dim, dim)
  86.         self.proj_drop = nn.Dropout(proj_drop_ratio)
  87.  
  88.     def forward(self, x):
  89.         # [batch_size, num_patches + 1, total_embed_dim]
  90.         B, N, C = x.shape
  91.  
  92.         # qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim]
  93.         # reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head]
  94.         # permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]
  95.         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
  96.         # [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
  97.         q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
  98.  
  99.         # transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1]
  100.         # @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1]
  101.         attn = (q @ k.transpose(-2, -1)) * self.scale
  102.         attn = attn.softmax(dim=-1)
  103.         attn = self.attn_drop(attn)
  104.  
  105.         # @: multiply -> [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
  106.         # transpose: -> [batch_size, num_patches + 1, num_heads, embed_dim_per_head]
  107.         # reshape: -> [batch_size, num_patches + 1, total_embed_dim]
  108.         x = (attn @ v).transpose(1, 2).reshape(B, N, C)
  109.         x = self.proj(x)
  110.         x = self.proj_drop(x)
  111.         return x
  112.  
  113.  
  114. class Mlp(nn.Module):
  115.     """
  116.     MLP as used in Vision Transformer, MLP-Mixer and related networks
  117.     """
  118.     def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
  119.         super().__init__()
  120.         out_features = out_features or in_features
  121.         hidden_features = hidden_features or in_features
  122.         self.fc1 = nn.Linear(in_features, hidden_features)
  123.         self.act = act_layer()
  124.         self.fc2 = nn.Linear(hidden_features, out_features)
  125.         self.drop = nn.Dropout(drop)
  126.  
  127.     def forward(self, x):
  128.         x = self.fc1(x)
  129.         x = self.act(x)
  130.         x = self.drop(x)
  131.         x = self.fc2(x)
  132.         x = self.drop(x)
  133.         return x
  134.  
  135.  
  136. class Block(nn.Module):
  137.     def __init__(self,
  138.                  dim,
  139.                  num_heads,
  140.                  mlp_ratio=4.,
  141.                  qkv_bias=False,
  142.                  qk_scale=None,
  143.                  drop_ratio=0.,
  144.                  attn_drop_ratio=0.,
  145.                  drop_path_ratio=0.,
  146.                  act_layer=nn.GELU,
  147.                  norm_layer=nn.LayerNorm):
  148.         super(Block, self).__init__()
  149.         self.norm1 = norm_layer(dim)
  150.         self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
  151.                               attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)
  152.         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
  153.         self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
  154.         self.norm2 = norm_layer(dim)
  155.         mlp_hidden_dim = int(dim * mlp_ratio)
  156.         self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)
  157.  
  158.     def forward(self, x):
  159.         x = x + self.drop_path(self.attn(self.norm1(x)))
  160.         x = x + self.drop_path(self.mlp(self.norm2(x)))
  161.         return x
  162.  
  163.  
  164. class VisionTransformer(nn.Module):
  165.     def __init__(self, img_size=224, patch_size=16, in_c=3, num_classes=1000,
  166.                  embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True,
  167.                  qk_scale=None, representation_size=None, distilled=False, drop_ratio=0.,
  168.                  attn_drop_ratio=0., drop_path_ratio=0., embed_layer=PatchEmbed, norm_layer=None,
  169.                  act_layer=None):
  170.         """
  171.         Args:
  172.             img_size (int, tuple): input image size
  173.             patch_size (int, tuple): patch size
  174.             in_c (int): number of input channels
  175.             num_classes (int): number of classes for classification head
  176.             embed_dim (int): embedding dimension
  177.             depth (int): depth of transformer
  178.             num_heads (int): number of attention heads
  179.             mlp_ratio (int): ratio of mlp hidden dim to embedding dim
  180.             qkv_bias (bool): enable bias for qkv if True
  181.             qk_scale (float): override default qk scale of head_dim ** -0.5 if set
  182.             representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
  183.             distilled (bool): model includes a distillation token and head as in DeiT models
  184.             drop_ratio (float): dropout rate
  185.             attn_drop_ratio (float): attention dropout rate
  186.             drop_path_ratio (float): stochastic depth rate
  187.             embed_layer (nn.Module): patch embedding layer
  188.             norm_layer: (nn.Module): normalization layer
  189.         """
  190.         super(VisionTransformer, self).__init__()
  191.         self.num_classes = num_classes
  192.         self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
  193.         self.num_tokens = 2 if distilled else 1
  194.         norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
  195.         act_layer = act_layer or nn.GELU
  196.  
  197.         self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_c=in_c, embed_dim=embed_dim)
  198.         num_patches = self.patch_embed.num_patches
  199.  
  200.         self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
  201.         self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
  202.         self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
  203.         self.pos_drop = nn.Dropout(p=drop_ratio)
  204.  
  205.         dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)]  # stochastic depth decay rule
  206.         self.blocks = nn.Sequential(*[
  207.             Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
  208.                   drop_ratio=drop_ratio, attn_drop_ratio=attn_drop_ratio, drop_path_ratio=dpr[i],
  209.                   norm_layer=norm_layer, act_layer=act_layer)
  210.             for i in range(depth)
  211.         ])
  212.         self.norm = norm_layer(embed_dim)
  213.  
  214.         # Representation layer
  215.         if representation_size and not distilled:
  216.             self.has_logits = True
  217.             self.num_features = representation_size
  218.             self.pre_logits = nn.Sequential(OrderedDict([
  219.                 ("fc", nn.Linear(embed_dim, representation_size)),
  220.                 ("act", nn.Tanh())
  221.             ]))
  222.         else:
  223.             self.has_logits = False
  224.             self.pre_logits = nn.Identity()
  225.  
  226.         # Classifier head(s)
  227.         self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
  228.         self.head_dist = None
  229.         if distilled:
  230.             self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()
  231.  
  232.         # Weight init
  233.         nn.init.trunc_normal_(self.pos_embed, std=0.02)
  234.         if self.dist_token is not None:
  235.             nn.init.trunc_normal_(self.dist_token, std=0.02)
  236.  
  237.         nn.init.trunc_normal_(self.cls_token, std=0.02)
  238.         self.apply(_init_vit_weights)
  239.  
  240.     def forward_features(self, x):
  241.         # [B, C, H, W] -> [B, num_patches, embed_dim]
  242.         x = self.patch_embed(x)  # [B, 196, 768]
  243.         # [1, 1, 768] -> [B, 1, 768]
  244.         cls_token = self.cls_token.expand(x.shape[0], -1, -1)
  245.         if self.dist_token is None:
  246.             x = torch.cat((cls_token, x), dim=1)  # [B, 197, 768]
  247.         else:
  248.             x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
  249.  
  250.         x = self.pos_drop(x + self.pos_embed)
  251.         x = self.blocks(x)
  252.         x = self.norm(x)
  253.         if self.dist_token is None:
  254.             return self.pre_logits(x[:, 0])
  255.         else:
  256.             return x[:, 0], x[:, 1]
  257.  
  258.     def forward(self, x):
  259.         x = self.forward_features(x)
  260.         if self.head_dist is not None:
  261.             x, x_dist = self.head(x[0]), self.head_dist(x[1])
  262.             if self.training and not torch.jit.is_scripting():
  263.                 # during inference, return the average of both classifier predictions
  264.                 return x, x_dist
  265.             else:
  266.                 return (x + x_dist) / 2
  267.         else:
  268.             x = self.head(x)
  269.         return x
  270.  
  271.  
  272. def _init_vit_weights(m):
  273.     """
  274.     ViT weight initialization
  275.     :param m: module
  276.     """
  277.     if isinstance(m, nn.Linear):
  278.         nn.init.trunc_normal_(m.weight, std=.01)
  279.         if m.bias is not None:
  280.             nn.init.zeros_(m.bias)
  281.     elif isinstance(m, nn.Conv2d):
  282.         nn.init.kaiming_normal_(m.weight, mode="fan_out")
  283.         if m.bias is not None:
  284.             nn.init.zeros_(m.bias)
  285.     elif isinstance(m, nn.LayerNorm):
  286.         nn.init.zeros_(m.bias)
  287.         nn.init.ones_(m.weight)
  288.  
  289.  
  290. def vit_base_patch16_224(num_classes: int = 1000):
  291.     """
  292.     ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
  293.     ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.
  294.     weights ported from official Google JAX impl:
  295.     链接: https://pan.baidu.com/s/1zqb08naP0RPqqfSXfkB2EA  密码: eu9f
  296.     """
  297.     model = VisionTransformer(img_size=224,
  298.                               patch_size=16,
  299.                               embed_dim=768,
  300.                               depth=12,
  301.                               num_heads=12,
  302.                               representation_size=None,
  303.                               num_classes=num_classes)
  304.     return model
  305.  
  306.  
  307. def vit_base_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = True):
  308.     """
  309.     ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
  310.     ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
  311.     weights ported from official Google JAX impl:
  312.     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch16_224_in21k-e5005f0a.pth
  313.     """
  314.     model = VisionTransformer(img_size=224,
  315.                               patch_size=16,
  316.                               embed_dim=768,
  317.                               depth=12,
  318.                               num_heads=12,
  319.                               representation_size=768 if has_logits else None,
  320.                               num_classes=num_classes)
  321.     return model
  322.  
  323.  
  324. def vit_base_patch32_224(num_classes: int = 1000):
  325.     """
  326.     ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
  327.     ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.
  328.     weights ported from official Google JAX impl:
  329.     链接: https://pan.baidu.com/s/1hCv0U8pQomwAtHBYc4hmZg  密码: s5hl
  330.     """
  331.     model = VisionTransformer(img_size=224,
  332.                               patch_size=32,
  333.                               embed_dim=768,
  334.                               depth=12,
  335.                               num_heads=12,
  336.                               representation_size=None,
  337.                               num_classes=num_classes)
  338.     return model
  339.  
  340.  
  341. def vit_base_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = True):
  342.     """
  343.     ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
  344.     ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
  345.     weights ported from official Google JAX impl:
  346.     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch32_224_in21k-8db57226.pth
  347.     """
  348.     model = VisionTransformer(img_size=224,
  349.                               patch_size=32,
  350.                               embed_dim=768,
  351.                               depth=12,
  352.                               num_heads=12,
  353.                               representation_size=768 if has_logits else None,
  354.                               num_classes=num_classes)
  355.     return model
  356.  
  357.  
  358. def vit_large_patch16_224(num_classes: int = 1000):
  359.     """
  360.     ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
  361.     ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.
  362.     weights ported from official Google JAX impl:
  363.     链接: https://pan.baidu.com/s/1cxBgZJJ6qUWPSBNcE4TdRQ  密码: qqt8
  364.     """
  365.     model = VisionTransformer(img_size=224,
  366.                               patch_size=16,
  367.                               embed_dim=1024,
  368.                               depth=24,
  369.                               num_heads=16,
  370.                               representation_size=None,
  371.                               num_classes=num_classes)
  372.     return model
  373.  
  374.  
  375. def vit_large_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = True):
  376.     """
  377.     ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
  378.     ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
  379.     weights ported from official Google JAX impl:
  380.     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch16_224_in21k-606da67d.pth
  381.     """
  382.     model = VisionTransformer(img_size=224,
  383.                               patch_size=16,
  384.                               embed_dim=1024,
  385.                               depth=24,
  386.                               num_heads=16,
  387.                               representation_size=1024 if has_logits else None,
  388.                               num_classes=num_classes)
  389.     return model
  390.  
  391.  
  392. def vit_large_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = True):
  393.     """
  394.     ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).
  395.     ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
  396.     weights ported from official Google JAX impl:
  397.     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth
  398.     """
  399.     model = VisionTransformer(img_size=224,
  400.                               patch_size=32,
  401.                               embed_dim=1024,
  402.                               depth=24,
  403.                               num_heads=16,
  404.                               representation_size=1024 if has_logits else None,
  405.                               num_classes=num_classes)
  406.     return model
  407.  
  408.  
  409. def vit_huge_patch14_224_in21k(num_classes: int = 21843, has_logits: bool = True):
  410.     """
  411.     ViT-Huge model (ViT-H/14) from original paper (https://arxiv.org/abs/2010.11929).
  412.     ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
  413.     NOTE: converted weights not currently available, too large for github release hosting.
  414.     """
  415.     model = VisionTransformer(img_size=224,
  416.                               patch_size=14,
  417.                               embed_dim=1280,
  418.                               depth=32,
  419.                               num_heads=16,
  420.                               representation_size=1280 if has_logits else None,
  421.                               num_classes=num_classes)
  422.     return model

2.utils.py

  1. import os
  2. import sys
  3. import json
  4. import pickle
  5. import random
  6.  
  7. import torch
  8. from tqdm import tqdm
  9.  
  10. import matplotlib.pyplot as plt
  11.  
  12.  
  13. def read_split_data(root: str, val_rate: float = 0.2):
  14.     random.seed(0)  # 保证随机结果可复现
  15.     assert os.path.exists(root), "dataset root: {} does not exist.".format(root)
  16.  
  17.     # 遍历文件夹，一个文件夹对应一个类别
  18.     flower_class = [cla for cla in os.listdir(root) if os.path.isdir(os.path.join(root, cla))]
  19.     # 排序，保证顺序一致
  20.     flower_class.sort()
  21.     # 生成类别名称以及对应的数字索引
  22.     class_indices = dict((k, v) for v, k in enumerate(flower_class))
  23.     json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4)
  24.     with open('class_indices.json', 'w') as json_file:
  25.         json_file.write(json_str)
  26.  
  27.     train_images_path = []  # 存储训练集的所有图片路径
  28.     train_images_label = []  # 存储训练集图片对应索引信息
  29.     val_images_path = []  # 存储验证集的所有图片路径
  30.     val_images_label = []  # 存储验证集图片对应索引信息
  31.     every_class_num = []  # 存储每个类别的样本总数
  32.     supported = [".jpg", ".JPG", ".png", ".PNG"]  # 支持的文件后缀类型
  33.     # 遍历每个文件夹下的文件
  34.     for cla in flower_class:
  35.         cla_path = os.path.join(root, cla)
  36.         # 遍历获取supported支持的所有文件路径
  37.         images = [os.path.join(root, cla, i) for i in os.listdir(cla_path)
  38.                   if os.path.splitext(i)[-1] in supported]
  39.         # 获取该类别对应的索引
  40.         image_class = class_indices[cla]
  41.         # 记录该类别的样本数量
  42.         every_class_num.append(len(images))
  43.         # 按比例随机采样验证样本
  44.         val_path = random.sample(images, k=int(len(images) * val_rate))
  45.  
  46.         for img_path in images:
  47.             if img_path in val_path:  # 如果该路径在采样的验证集样本中则存入验证集
  48.                 val_images_path.append(img_path)
  49.                 val_images_label.append(image_class)
  50.             else:  # 否则存入训练集
  51.                 train_images_path.append(img_path)
  52.                 train_images_label.append(image_class)
  53.  
  54.     print("{} images were found in the dataset.".format(sum(every_class_num)))
  55.     print("{} images for training.".format(len(train_images_path)))
  56.     print("{} images for validation.".format(len(val_images_path)))
  57.  
  58.     plot_image = False
  59.     if plot_image:
  60.         # 绘制每种类别个数柱状图
  61.         plt.bar(range(len(flower_class)), every_class_num, align='center')
  62.         # 将横坐标0,1,2,3,4替换为相应的类别名称
  63.         plt.xticks(range(len(flower_class)), flower_class)
  64.         # 在柱状图上添加数值标签
  65.         for i, v in enumerate(every_class_num):
  66.             plt.text(x=i, y=v + 5, s=str(v), ha='center')
  67.         # 设置x坐标
  68.         plt.xlabel('image class')
  69.         # 设置y坐标
  70.         plt.ylabel('number of images')
  71.         # 设置柱状图的标题
  72.         plt.title('flower class distribution')
  73.         plt.show()
  74.  
  75.     return train_images_path, train_images_label, val_images_path, val_images_label
  76.  
  77.  
  78. def plot_data_loader_image(data_loader):
  79.     batch_size = data_loader.batch_size
  80.     plot_num = min(batch_size, 4)
  81.  
  82.     json_path = './class_indices.json'
  83.     assert os.path.exists(json_path), json_path + " does not exist."
  84.     json_file = open(json_path, 'r')
  85.     class_indices = json.load(json_file)
  86.  
  87.     for data in data_loader:
  88.         images, labels = data
  89.         for i in range(plot_num):
  90.             # [C, H, W] -> [H, W, C]
  91.             img = images[i].numpy().transpose(1, 2, 0)
  92.             # 反Normalize操作
  93.             img = (img * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]) * 255
  94.             label = labels[i].item()
  95.             plt.subplot(1, plot_num, i+1)
  96.             plt.xlabel(class_indices[str(label)])
  97.             plt.xticks([])  # 去掉x轴的刻度
  98.             plt.yticks([])  # 去掉y轴的刻度
  99.             plt.imshow(img.astype('uint8'))
  100.         plt.show()
  101.  
  102.  
  103. def write_pickle(list_info: list, file_name: str):
  104.     with open(file_name, 'wb') as f:
  105.         pickle.dump(list_info, f)
  106.  
  107.  
  108. def read_pickle(file_name: str) -> list:
  109.     with open(file_name, 'rb') as f:
  110.         info_list = pickle.load(f)
  111.         return info_list
  112.  
  113.  
  114. def train_one_epoch(model, optimizer, data_loader, device, epoch):
  115.     model.train()
  116.     loss_function = torch.nn.CrossEntropyLoss()
  117.     accu_loss = torch.zeros(1).to(device)  # 累计损失
  118.     accu_num = torch.zeros(1).to(device)   # 累计预测正确的样本数
  119.     optimizer.zero_grad()
  120.  
  121.     sample_num = 0
  122.     data_loader = tqdm(data_loader, file=sys.stdout)
  123.     for step, data in enumerate(data_loader):
  124.         images, labels = data
  125.         sample_num += images.shape[0]
  126.  
  127.         pred = model(images.to(device))
  128.         pred_classes = torch.max(pred, dim=1)[1]
  129.         accu_num += torch.eq(pred_classes, labels.to(device)).sum()
  130.  
  131.         loss = loss_function(pred, labels.to(device))
  132.         loss.backward()
  133.         accu_loss += loss.detach()
  134.  
  135.         data_loader.desc = "[train epoch {}] loss: {:.3f}, acc: {:.3f}".format(epoch,
  136.                                                                                accu_loss.item() / (step + 1),
  137.                                                                                accu_num.item() / sample_num)
  138.  
  139.         if not torch.isfinite(loss):
  140.             print('WARNING: non-finite loss, ending training ', loss)
  141.             sys.exit(1)
  142.  
  143.         optimizer.step()
  144.         optimizer.zero_grad()
  145.  
  146.     return accu_loss.item() / (step + 1), accu_num.item() / sample_num
  147.  
  148.  
  149. @torch.no_grad()
  150. def evaluate(model, data_loader, device, epoch):
  151.     loss_function = torch.nn.CrossEntropyLoss()
  152.  
  153.     model.eval()
  154.  
  155.     accu_num = torch.zeros(1).to(device)   # 累计预测正确的样本数
  156.     accu_loss = torch.zeros(1).to(device)  # 累计损失
  157.  
  158.     sample_num = 0
  159.     data_loader = tqdm(data_loader, file=sys.stdout)
  160.     for step, data in enumerate(data_loader):
  161.         images, labels = data
  162.         sample_num += images.shape[0]
  163.  
  164.         pred = model(images.to(device))
  165.         pred_classes = torch.max(pred, dim=1)[1]
  166.         accu_num += torch.eq(pred_classes, labels.to(device)).sum()
  167.  
  168.         loss = loss_function(pred, labels.to(device))
  169.         accu_loss += loss
  170.  
  171.         data_loader.desc = "[valid epoch {}] loss: {:.3f}, acc: {:.3f}".format(epoch,
  172.                                                                                accu_loss.item() / (step + 1),
  173.                                                                                accu_num.item() / sample_num)
  174.  
  175.     return accu_loss.item() / (step + 1), accu_num.item() / sample_num

3.my_dataset.py

  1. from PIL import Image
  2. import torch
  3. from torch.utils.data import Dataset
  4.  
  5.  
  6. class MyDataSet(Dataset):
  7.     """自定义数据集"""
  8.  
  9.     def __init__(self, images_path: list, images_class: list, transform=None):
  10.         self.images_path = images_path
  11.         self.images_class = images_class
  12.         self.transform = transform
  13.  
  14.     def __len__(self):
  15.         return len(self.images_path)
  16.  
  17.     def __getitem__(self, item):
  18.         img = Image.open(self.images_path[item])
  19.         # RGB为彩色图片，L为灰度图片
  20.         if img.mode != 'RGB':
  21.             raise ValueError("image: {} isn't RGB mode.".format(self.images_path[item]))
  22.         label = self.images_class[item]
  23.  
  24.         if self.transform is not None:
  25.             img = self.transform(img)
  26.  
  27.         return img, label
  28.  
  29.     @staticmethod
  30.     def collate_fn(batch):
  31.         # 官方实现的default_collate可以参考
  32.         # https://github.com/pytorch/pytorch/blob/67b7e751e6b5931a9f45274653f4f653a4e6cdf6/torch/utils/data/_utils/collate.py
  33.         images, labels = tuple(zip(*batch))
  34.  
  35.         images = torch.stack(images, dim=0)
  36.         labels = torch.as_tensor(labels)
  37.         return images, labels

4.train.py

其中若使用预训练模型需要提前下载，下载地址在 utils.py 处有标明，代码默认是使用预训练模型的。下载后，预训练模型放入项目的根目录即可。我训练的数据集种类有三种，于是我将网络的全连接层的输出改成了 3 ，各位需要依据自己数据集不同来进行调整。

若下载不方便，也可以下载我上传的资源：

vit_base_patch16_224_in21k.zip-深度学习文档类资源-CSDN下载

  1. import os
  2. import math
  3. import argparse
  4.  
  5. import torch
  6. import torch.optim as optim
  7. import torch.optim.lr_scheduler as lr_scheduler
  8. from torch.utils.tensorboard import SummaryWriter
  9. from torchvision import transforms
  10.  
  11.  
  12. from my_dataset import MyDataSet
  13. from vit_model import vit_base_patch16_224_in21k as create_model
  14. from utils import read_split_data, train_one_epoch, evaluate
  15.  
  16. import xlwt
  17.  
  18. book = xlwt.Workbook(encoding='utf-8') #创建Workbook，相当于创建Excel
  19. # 创建sheet，Sheet1为表的名字，cell_overwrite_ok为是否覆盖单元格
  20. sheet1 = book.add_sheet(u'Train_data', cell_overwrite_ok=True)
  21. # 向表中添加数据
  22. sheet1.write(0, 0, 'epoch')
  23. sheet1.write(0, 1, 'Train_Loss')
  24. sheet1.write(0, 2, 'Train_Acc')
  25. sheet1.write(0, 3, 'Val_Loss')
  26. sheet1.write(0, 4, 'Val_Acc')
  27. sheet1.write(0, 5, 'lr')
  28. sheet1.write(0, 6, 'Best val Acc')
  29.  
  30. def main(args):
  31.     best_acc = 0
  32.  
  33.     device = torch.device(args.device if torch.cuda.is_available() else "cpu")
  34.  
  35.     if os.path.exists("./weights") is False:
  36.         os.makedirs("./weights")
  37.  
  38.     tb_writer = SummaryWriter()
  39.  
  40.     train_images_path, train_images_label, val_images_path, val_images_label = read_split_data(args.data_path)
  41.  
  42.     data_transform = {
  43.         "train": transforms.Compose([transforms.RandomResizedCrop(224),
  44.                                      transforms.RandomHorizontalFlip(),
  45.                                      transforms.ToTensor(),
  46.                                      transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
  47.         "val": transforms.Compose([transforms.Resize(256),
  48.                                    transforms.CenterCrop(224),
  49.                                    transforms.ToTensor(),
  50.                                    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])}
  51.  
  52.     # 实例化训练数据集
  53.     train_dataset = MyDataSet(images_path=train_images_path,
  54.                               images_class=train_images_label,
  55.                               transform=data_transform["train"])
  56.  
  57.     # 实例化验证数据集
  58.     val_dataset = MyDataSet(images_path=val_images_path,
  59.                             images_class=val_images_label,
  60.                             transform=data_transform["val"])
  61.  
  62.     batch_size = args.batch_size
  63.     nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
  64.     print('Using {} dataloader workers every process'.format(nw))
  65.     train_loader = torch.utils.data.DataLoader(train_dataset,
  66.                                                batch_size=batch_size,
  67.                                                shuffle=True,
  68.                                                pin_memory=True,
  69.                                                num_workers=nw,
  70.                                                collate_fn=train_dataset.collate_fn)
  71.  
  72.     val_loader = torch.utils.data.DataLoader(val_dataset,
  73.                                              batch_size=batch_size,
  74.                                              shuffle=False,
  75.                                              pin_memory=True,
  76.                                              num_workers=nw,
  77.                                              collate_fn=val_dataset.collate_fn)
  78.  
  79.     model = create_model(num_classes=3, has_logits=False).to(device)
  80.  
  81.     images = torch.zeros(1, 3, 224, 224).to(device)#要求大小与输入图片的大小一致
  82.     tb_writer.add_graph(model, images, verbose=False)
  83.  
  84.     if args.weights != "":
  85.         assert os.path.exists(args.weights), "weights file: '{}' not exist.".format(args.weights)
  86.         weights_dict = torch.load(args.weights, map_location=device)
  87.         # 删除不需要的权重
  88.         del_keys = ['head.weight', 'head.bias'] if model.has_logits \
  89.             else ['pre_logits.fc.weight', 'pre_logits.fc.bias', 'head.weight', 'head.bias']
  90.         for k in del_keys:
  91.             del weights_dict[k]
  92.         print(model.load_state_dict(weights_dict, strict=False))
  93.  
  94.     if args.freeze_layers:
  95.         for name, para in model.named_parameters():
  96.             # 除head, pre_logits外，其他权重全部冻结
  97.             if "head" not in name and "pre_logits" not in name:
  98.                 para.requires_grad_(False)
  99.             else:
  100.                 print("training {}".format(name))
  101.  
  102.     pg = [p for p in model.parameters() if p.requires_grad]
  103.  
  104.     optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=5E-5)
  105.     # Scheduler https://arxiv.org/pdf/1812.01187.pdf
  106.     lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf  # cosine
  107.     scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
  108.  
  109.     for epoch in range(args.epochs):
  110.  
  111.         sheet1.write(epoch+1, 0, epoch+1)
  112.         sheet1.write(epoch + 1, 5, str(optimizer.state_dict()['param_groups'][0]['lr']))
  113.  
  114.         # train
  115.         train_loss, train_acc = train_one_epoch(model=model,
  116.                                                 optimizer=optimizer,
  117.                                                 data_loader=train_loader,
  118.                                                 device=device,
  119.                                                 epoch=epoch)
  120.  
  121.         scheduler.step()
  122.  
  123.         sheet1.write(epoch + 1, 1, str(train_loss))
  124.         sheet1.write(epoch + 1, 2, str(train_acc))
  125.  
  126.         # validate
  127.         val_loss, val_acc = evaluate(model=model,
  128.                                      data_loader=val_loader,
  129.                                      device=device,
  130.                                      epoch=epoch)
  131.  
  132.         sheet1.write(epoch + 1, 3, str(val_loss))
  133.         sheet1.write(epoch + 1, 4, str(val_acc))
  134.  
  135.         tags = ["train_loss", "train_acc", "val_loss", "val_acc", "learning_rate"]
  136.         tb_writer.add_scalar(tags[0], train_loss, epoch)
  137.         tb_writer.add_scalar(tags[1], train_acc, epoch)
  138.         tb_writer.add_scalar(tags[2], val_loss, epoch)
  139.         tb_writer.add_scalar(tags[3], val_acc, epoch)
  140.         tb_writer.add_scalar(tags[4], optimizer.param_groups[0]["lr"], epoch)
  141.  
  142.  
  143.         if val_acc > best_acc:
  144.             best_acc = val_acc
  145.             torch.save(model.state_dict(), "./weights/best_model.pth")
  146.             #torch.save(model.state_dict(), "./weights/model-{}.pth".format(epoch))
  147.  
  148.     sheet1.write(1, 6, str(best_acc))
  149.     book.save('.\Train_data.xlsx')
  150.     print("The Best Acc = : {:.4f}".format(best_acc))
  151.  
  152.  
  153.  
  154. if __name__ == '__main__':
  155.     parser = argparse.ArgumentParser()
  156.     parser.add_argument('--num_classes', type=int, default=3)
  157.     parser.add_argument('--epochs', type=int, default=100)
  158.     parser.add_argument('--batch-size', type=int, default=8)
  159.     parser.add_argument('--lr', type=float, default=0.001)
  160.     parser.add_argument('--lrf', type=float, default=0.01)
  161.  
  162.     # 数据集所在根目录
  163.     parser.add_argument('--data-path', type=str,
  164.                         default=r"D:\pyCharmdata\resnet50_plant_3\datasets\train")
  165.     parser.add_argument('--model-name', default='', help='create model name')
  166.  
  167.     # 预训练权重路径，如果不想载入就设置为空字符
  168.     parser.add_argument('--weights', type=str, default='./vit_base_patch16_224_in21k.pth',
  169.                         help='initial weights path')
  170.  
  171.     # 是否冻结权重
  172.     parser.add_argument('--freeze-layers', type=bool, default=False)
  173.     parser.add_argument('--device', default='cuda:0', help='device id (i.e. 0 or 0,1 or cpu)')
  174.  
  175.     opt = parser.parse_args()
  176.  
  177.     main(opt)

5.predict.py

可以实现单张图片的种类预测，得分最高的便是模型预测种类。

  1. import os
  2. import json
  3.  
  4. import torch
  5. from PIL import Image
  6. from torchvision import transforms
  7. import matplotlib.pyplot as plt
  8.  
  9. from vit_model import vit_base_patch16_224_in21k as create_model
  10.  
  11.  
  12. def main():
  13.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  14.  
  15.     data_transform = transforms.Compose(
  16.         [transforms.Resize(256),
  17.          transforms.CenterCrop(224),
  18.          transforms.ToTensor(),
  19.          transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
  20.  
  21.     # load image
  22.     img_path = r"D:\pyCharmdata\resnet50_plant_3\datasets\test\Huanglong_disease\000000.jpg"
  23.     assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
  24.     img = Image.open(img_path)
  25.     plt.imshow(img)
  26.     # [N, C, H, W]
  27.     img = data_transform(img)
  28.     # expand batch dimension
  29.     img = torch.unsqueeze(img, dim=0)
  30.  
  31.     # read class_indict
  32.     json_path = './class_indices.json'
  33.     assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
  34.  
  35.     with open(json_path, "r") as f:
  36.         class_indict = json.load(f)
  37.  
  38.     # create model
  39.     model = create_model(num_classes=3, has_logits=False).to(device)
  40.     # load model weights
  41.     model_weight_path = "./weights/best_model.pth"
  42.     model.load_state_dict(torch.load(model_weight_path, map_location=device))
  43.     model.eval()
  44.     with torch.no_grad():
  45.         # predict class
  46.         output = torch.squeeze(model(img.to(device))).cpu()
  47.         predict = torch.softmax(output, dim=0)
  48.         predict_cla = torch.argmax(predict).numpy()
  49.  
  50.     print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
  51.                                                  predict[predict_cla].numpy())
  52.     plt.title(print_res)
  53.     for i in range(len(predict)):
  54.         print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
  55.                                                   predict[i].numpy()))
  56.     plt.show()
  57.  
  58.  
  59. if __name__ == '__main__':
  60.     main()

预测结果展示：

四、训练数据

在配置好环境和数据集、预训练模型的路径后，即可运行 train.py 开始训练，默认是训练100轮。

训练使用的是SGDM优化器，初始学习率为0.001，使用LambdaLR自定义学习率调整策略，导入预训练模型但不冻结网络层和参数。

 训练过程中可以在项目路径下的终端 输入：

tensorboard --logdir=runs/

进行实时监控训练进程，也可以查看 Vision Transformer 的网络可视化结构。

Vision Transformer 的网络可视化 ：

我简单训练了100轮后，最高 val_acc 准确率为 0.9976。

 训练结束后，会在项目根目录生成一个Excel文件，里面记载了训练全过程的数据，你也可以在通过 Matlab 来获得高度自定义化的可视化对比图片，堪称 论文人 的福音。

我这里只展示前10轮的训练数据。

我的完整项目框架，有需要的自取：

Vit_myself.zip-深度学习文档类资源-CSDN下载

 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

如果本文对你有帮助，欢迎一键三连！！！

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/weixin_40725706/article/detail/359149
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号