大家好，今天和各位分享一下如何使用 Pytorch 构建 Vision Transformer 网络模型，并使用 权重迁移学习方法 训练模型并预测。

Vision Transformer 的原理和 TensorFlow2 的实现方法可以看我下面这篇博文：https://blog.csdn.net/dgvv4/article/details/124792386

1. 引言

经典的 Transformer 由 Encoder 和 Decoder 组成，其中，最重要的就是多头注意力机制（Multi-head attention）。在 Vision Transformer 中，作者通过 Transformer 的 Encoder 部分直接进行分类任务，与 NLP 中的处理方式类似，在图片序列中加入分类 token，图片序列则由原始图像切割成若干个图像块（Patch）得到，如下图所示。

主要通过以下步骤进行转换：

（1）假设一张输入图片的形状为（B,C,H,W），其中 B 代表 Batch 大小，C 表示输入图片通道个数，H 和 W 表示输入图片的高和宽。

（2）那么，通过矩阵变换，可以将其转化为 $(B,N,P^{2}C)$ ，其中，P 代表 Patch 的大小，N 的值为 $H*W/P^{2}$ 。

（3）将切分为若干个 Patch 的图片分别送到 TransFormer Layer 中处理，在此过程中通过注意力机制进行输入特征的提取。

2. 模型构建

接下来构建 Vision Transformer 的主干模型，本小节的代码都写在 VisionTransformer_model.py 文件中。先导入模型构建过程中需要用到的工具包。

import torch
from torch import nn
from functools import partial

2.1 Patch Embedding

首先对输入图像 [b,3,224,224] 做Patch Embedding 操作。首先进行图像分块，将图片切分成14*14个图像块（Patch），每个 Patch 的尺寸为 16*16。通过提取输入图片中的平坦像素向量，将每个输入 Patch 送入线性投影层，得到 Patch Embeddings。

在代码中其流程如上图，先经过一个 kernel=(16,16)，strides=16 的卷积层划分图像块，再将 h和w 维度整合为 num_patches 维度，代表一共有 196 个 patch，每个 patch 为 16*16

代码如下：

# --------------------------------------- #
#（1）patch embedding
'''
img_size=224 : 输入图像的宽高
patch_size=16 ： 每个patch的宽高，也是卷积核的尺寸和步长
in_c=3 ： 输入图像的通道数
embed_dim=768 ： 卷积输出通道数
'''
# --------------------------------------- #
class patchembed(nn.Module):# 初始化def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768):super(patchembed, self).__init__()# 输入图像的尺寸224*224self.img_size = (img_size, img_size)# 每个patch的大小16*16self.patch_size = (patch_size, patch_size)# 将输入图像划分成14*14个patchself.grid_size = (img_size//patch_size, img_size//patch_size)# 一共有14*14个patchself.num_patches = self.grid_size[0] * self.grid_size[1]# 使用16*16的卷积切分图像，将图像分成14*14个self.proj = nn.Conv2d(in_channels=in_c, out_channels=embed_dim, kernel_size=patch_size, stride=patch_size)# 定义标准化方法，给LN传入默认参数epsnorm_layer = partial(nn.LayerNorm, eps=1e-6)self.norm = norm_layer(embed_dim)# 前向传播def forward(self, inputs):# 获得输入图像的shapeB, C, H, W = inputs.shape# 如果输入图像的宽高不等于224*224就报错assert H==self.img_size[0] and W==self.img_size[1], 'input shape does not match 224*224'# 卷积层切分patch [b,3,224,224]==>[b,768,14,14]x = self.proj(inputs)# 展平 [b,768,14,14]==>[b,768,14*14]x = x.flatten(start_dim=2, end_dim=-1)  # 将索引为 start_dim 和 end_dim 之间（包括该位置）的数量相乘# 维度调整 [b,768,14*14]==>[b,14*14,768]x = x.transpose(1, 2)  # 实现一个张量的两个轴之间的维度转换# 标准化x = self.norm(x)return x

2.2 类别标签和位置编码

为了输出融合了全局语义信息的向量表示，在第一个输入张量前添加可学习分类变量。经过编码器编码后，在最后一层输出中，该位置对应的输出张量就可以用于分类任务。与其他位置对应的输出向量相比，该向量可以更好的融合图像中各个图像块之间的依赖关系。

在 Transformer 更新的过程中，输入序列的顺序信息会丢失。Transformer 本身并没有办法学习这个信息，所以需要一种方法将位置表示聚合到模型的输入嵌入中。我们对每个 Patch 进行位置编码，该位置编码采用随机初始化，之后参与模型训练。与传统三角函数的位置编码方法不同，该方法是可学习的。

最后，将 Patch-Embeddings 和 class-token 进行堆叠，和 Position-Embeddings 进行叠加，得到最终嵌入向量，该向量输入给 Transformer 层进行后续处理。

在代码中，要注意 cls_token 和 inputs 做堆叠 torch.cat() 时，需要将类别标签 cls_token 放在最前面。

代码如下：

# --------------------------------------- #
#（2）类别标签和位置标签
'''
embed_dim : 代表patchembed层输出的通道数
'''
# --------------------------------------- #
class class_token_pos_embed(nn.Module):# 初始化def __init__(self, embed_dim):super(class_token_pos_embed, self).__init__()# patchembed层将图像划分的patch个数==14*14num_patches = patchembed().num_patchesself.num_tokens = 1  # 类别标签# 创建可学习的类别标签 [1,1,768]self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))# 创建可学习的位置编码 [1,196+1,768]self.pos_embed = nn.Parameter(torch.zeros(1, num_patches+self.num_tokens, embed_dim))# 权重以正态分布初始化nn.init.trunc_normal_(self.pos_embed, std=0.02)nn.init.trunc_normal_(self.cls_token, std=0.02)# 前向传播def forward(self, x):  # 输入特征图的shape=[b,196,768]# 类别标签扩充维度 [1,1,768]==>[b,1,768]cls_token = self.cls_token.expand(x.shape[0], -1, -1)# 将类别标签添加到特征图中 [b,1,768]+[b,196,768]==>[b,197,768]x = torch.cat((cls_token, x), dim=1)# 添加位置编码 [b,197,768]+[1,197,768]==>[b,197,768]x = x + self.pos_embedreturn x

2.3 多头自注意力模块

Transformer 层中，主要包含多头注意力机制和多层感知机模块，下面先介绍多头自注意力模块。

单个的注意力机制，其每个输入包含三个不同的向量，分别为 Query向量（Q），Key向量（K），Value向量（V）。他们的结果分别由输入特征图和三个权重做矩阵乘法得到。

$q=X*W^{q} \;\;\;\;\;\ k=X*W^{k} \;\;\;\;\;\ v=X*W^{v}$

接着为每一个输入计算一个得分 $Score = q*k$

为了使梯度稳定，对 Score 的值进行归一化处理，并将结果通过 softmax 函数进行映射。之后再和 v 做矩阵相乘，得到加权后每个输入向量的得分 v。计算完后再乘以一个权重张量 W 提取特征。

计算公式如下，其中 $\sqrt{d_{k}}$ 代表 K 向量维度的平方根

$attention = Softmax(\frac{Q*K^{T}}{\sqrt{d_{k}}}*V)$

代码如下：

# --------------------------------------- #
#（3）多头注意力模块
'''
dim : 代表输入特征图的通道数
num_heads : 多头注意力中heads的个数
qkv_bias ： 生成qkv时是否使用偏置
atten_drop_ratio ：qk计算完之后的dropout层
proj_drop_ratio ： qkv计算完成之后的dropout层
'''
# --------------------------------------- #
class attention(nn.Module):# 初始化def __init__(self, dim, num_heads=12, qkv_bias=False, atten_drop_ratio=0., proj_drop_ratio=0.):super(attention, self).__init__()# 多头注意力的数量self.num_heads = num_heads  # 将生成的qkv均分成num_heads个。得到每个head的qkv对应的通道数。head_dim = dim // num_heads# 公式中的分母self.scale = head_dim ** -0.5# 通过一个全连接层计算qkvself.qkv = nn.Linear(in_features=dim, out_features=dim*3, bias=qkv_bias)# dropout层self.atten_drop = nn.Dropout(atten_drop_ratio)# 再qkv计算完之后通过一个全连接提取特征self.proj = nn.Linear(in_features=dim, out_features=dim)# dropout层self.proj_drop = nn.Dropout(proj_drop_ratio)# 前向传播def forward(self, inputs):# 获取输入图像的shape=[b,197,768]B, N, C = inputs.shape# 将输入特征图经过全连接层生成qkv [b,197,768]==>[b,197,768*3]qkv = self.qkv(inputs)# 维度调整 [b,197,768*3]==>[b, 197, 3, 12, 768//12]qkv = qkv.reshape(B, N, 3, self.num_heads, C//self.num_heads)# 维度重排==> [3, B, 12, 197, 768//12]qkv = qkv.permute(2,0,3,1,4)# 切片提取q、k、v的值，单个的shape=[B, 12, 197, 768//12]q, k, v = qkv[0], qkv[1], qkv[2]# 针对每个head计算 ==> [B, 12, 197, 197] atten = (q @ k.transpose(-2,-1)) * self.scale  # @ 代表在多维tensor的最后两个维度矩阵相乘# 对计算结果的每一行经过softmaxatten = atten.softmax(dim=-1)# dropout层atten = self.atten_drop(atten)# softmax后的结果和v加权 ==> [B, 12, 197, 768//12]x = atten @ v# 通道重排 ==> [B, 197, 12, 768//12]x = x.transpose(1,2)# 维度调整 ==> [B, 197, 768]x = x.reshape(B,N,C)# 通过全连接层融合特征 ==> [B, 197, 768]x = self.proj(x)# dropout层x = self.proj_drop(x)return x

2.4 MLP 多层感知器

这个部分简单来看就是两个全连接层提取特征，流程图如下。第一个全连接层通道上升4倍，第二个全连接层通道下降为原来。

代码如下：

# --------------------------------------- #
#（4）MLP多层感知器
'''
in_features : 输入特征图的通道数
hidden_features : 第一个全连接层上升通道数
out_features : 第二个全连接层的下降的通道数
drop : 全连接层后面的dropout层的杀死神经元的概率
'''
# --------------------------------------- #
class MLP(nn.Module):# 初始化def __init__(self, in_features, hidden_features, out_features=None, drop=0.):super(MLP, self).__init__()# MLP的输出通道数默认等于输入通道数out_features = out_features or in_features# 第一个全连接层上升通道数self.fc1 = nn.Linear(in_features=in_features, out_features=hidden_features)# GeLU激活函数self.act = nn.GELU()# 第二个全连接下降通道数self.fc2 = nn.Linear(in_features=hidden_features, out_features=out_features)# dropout层self.drop = nn.Dropout(drop)# 前向传播def forward(self, inputs):# [b,197,768]==>[b,197,3072]x = self.fc1(inputs)x = self.act(x)x = self.drop(x)# [b,197,3072]==>[b,197,768]x = self.fc2(x)x = self.drop(x)return x

2.5 特征提取模块

Transformer 的单个特征提取模块是由 多头注意力机制 和 多层感知机模块 组合而成，encoder_block 模块的流程图如下。

输入图像像经过 LayerNormalization 标准化后，再经过我们上面定义的多头注意力模块，将输出结果和输入特征图残差连接，图像在特征提取过程中shape保持不变。

将输出结果再经过标准化，然后送入多层感知器提取特征，再使用残差连接输入和输出。

而 transformer 的特征提取模块是由多个 encoder_block 叠加而成，这里连续使用12个 encoder_block 模块来提取特征。

代码如下：

# --------------------------------------- #
#（5）Encoder Block
'''
dim : 该模块的输入特征图个数
mlp_ratio ： MLP中第一个全连接层上升的通道数
drop_ratio : 该模块的dropout层的杀死神经元的概率
'''
# --------------------------------------- #
class encoder_block(nn.Module):# 初始化def __init__(self, dim, mlp_ratio=4., drop_ratio=0.):super(encoder_block, self).__init__()# LayerNormalization层self.norm1 = nn.LayerNorm(dim)# 实例化多头注意力self.atten = attention(dim)# dropoutself.drop = nn.Dropout()# LayerNormalization层self.norm2 = nn.LayerNorm(dim)# MLP中第一个全连接层上升的通道数hidden_features = int(dim * mlp_ratio)# MLP多层感知器self.mlp = MLP(in_features=dim, hidden_features=hidden_features)# 前向传播def forward(self, inputs):# [b,197,768]==>[b,197,768]x = self.norm1(inputs)x = self.atten(x)x = self.drop(x)feat1 = x + inputs  # 残差连接# [b,197,768]==>[b,197,768]x = self.norm2(feat1)x = self.mlp(x)x = self.drop(x)feat2 = x + feat1  # 残差连接return feat2

2.6 主干网络

接下来就搭建网络了，将上面所有的模块组合到一起，如下图所示。

在下面代码中要注意的是 x= x[:,0] 取出所有的类别标签。因为在 cls_pos_embed 模块中，我们将 cls_token 和输入图像在 patch 维度上堆叠，用于学习每张特征图的类别信息。最后经过一个全连接层得出每张图片属于每个类别的得分。

代码如下：

# --------------------------------------- #
#（6）主干网络
'''
num_class: 分类数
depth : 重复堆叠encoder_block的次数
drop_ratio : 位置编码后的dropout层
embed_dim : patchembed层输出通道数
'''
# --------------------------------------- #
class VIT(nn.Module):# 初始化def __init__(self, num_classes=1000, depth=12, drop_ratio=0., embed_dim=768):super(VIT, self).__init__()self.num_classes = num_classes  # 分类类别数# 实例化patchembed层self.patchembed = patchembed()# 实例化类别标签和位置编码self.cls_pos_embed = class_token_pos_embed(embed_dim=embed_dim)        # 位置编码后做dropoutself.pos_drop = nn.Dropout(drop_ratio)# 在列表中添加12个encoder_blockself.blocks = nn.Sequential(*[encoder_block(dim=embed_dim) for _ in range(depth)])# 定义LayerNormalization标准化方法norm_layer = partial(nn.LayerNorm, eps=1e-6)# 经过12个encoder之后的标准化层self.norm = norm_layer(embed_dim)# 分类层self.head = nn.Linear(in_features=embed_dim, out_features=num_classes)# 权值初始化for m in self.modules():# 对卷积层使用kaiming初始化if isinstance(m, nn.Conv2d):nn.init.kaiming_normal_(m.weight, mode='fan_out')# 对偏置初始化if m.bias is not None:nn.init.zeros_(m.bias)# 对标准化层初始化elif isinstance(m, nn.LayerNorm):nn.init.ones_(m.weight)nn.init.zeros_(m.bias)# 对全连接层初始化elif isinstance(m, nn.Linear):nn.init.normal_(m.weight, 0, 0.01)if m.bias is not None:nn.init.zeros_(m.bias)# 前向传播def forward(self, inputs):# 先将输入传递给patchembed [b,3,224,224]==>[b,196,768]x = self.patchembed(inputs)# 对特征图添加类别标签和位置编码x = self.cls_pos_embed(x)# dropout层x = self.pos_drop(x)# 经过12个encoder层==>[b,197,768]x = self.blocks(x)# LN标准化层x = self.norm(x)# 提取类别标签的输出,因为在cat时将类别标签放在最前面x = x[:, 0]  # [b,197,768]==>[b,768]# 全连接层分类 [b,768]==>[b,1000]x = self.head(x)return x

3. 训练阶段

接下来对使用权重迁移学习的方法训练模型，这里用的网络是 VIT B-16 模型，patch的尺寸为16*16，patchembedding的输出通道数为768。首先导入所有的工具包，定义好所有需要的参数，找到文件路径，方便后期使用管理。

import  torch
from torch import nn, optim
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from VisionTransformer_model import VIT  # 导入我们之前定义的 VIT B-16 模型
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 绘图显示中文# --------------------------------------------- #
#（0）参数设置
# --------------------------------------------- #
batch_size = 16  # 每个step处理16张图片
epochs = 10  # 训练10轮
best_loss = 1.0  # 当验证集损失小于1时才保存权重# 数据集目录位置
filepath = 'D:/deeplearning/test/数据集/animal/'
# 预训练权重位置
weightpath = 'D:/deeplearning/imgnet/pytorchimgnet/pretrained_weights/vit_base_patch16_224.pth'
# 训练时保存权重文件的位置
savepath = 'D:/deeplearning/imgnet/pytorchimgnet/save_weights/'# 获取GPU设备，检测到了就用GPU，检测不到就用CPU
if torch.cuda.is_available():device = torch.device('cuda:0')
else:device = torch.device('cpu')

3.1 构造数据集

首先定义训练集和验证集的数据预处理方法 data_transform。通过 transforms.Resize() 将输入图像的尺寸缩放到模型要求的 224*224 大小，然后再通过 transforms.ToTensor() 将像素值类型从 numpy 变成 tensor 类型，并归一化处理，像素值大小从 [0,255] 变换到 [0,1]，再调整输入图像的维度，从 [h,w,c] 变成 [c,h,w]；接着 transforms.Normalize() 对图像的每个颜色通道做标准化处理，使像素值满足正态分布。

预处理之后就构造训练集和验证集 dataloader，指定 batch_size=16，代表训练时每个 step 训练16张图片。

接着查看数据集信息，查看分类类别及其对应的索引信息，其中 datasets['train'].class_to_idx 的结果是 {'cats':0, 'dogs':1, 'panda':2}

代码如下：

# --------------------------------------------- #
#（1）数据集处理
# --------------------------------------------- #
# 定义预处理方法
data_transform = {# 训练集预处理方法'train' : transforms.Compose([transforms.Resize((224,224)),  # 将原始图片缩放至224*224大小transforms.RandomHorizontalFlip(),  # 随机水平翻转transforms.ToTensor(),  # numpy类型变tensor，维度调整，数据归一化transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])   # 对图像的三个通道分别做标准化]),# 验证集预处理方法'val' : transforms.Compose([transforms.Resize((224,224)),  # 将输入图像缩放至224*224大小transforms.ToTensor(),transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])])}# 加载数据集
datasets = {'train' : datasets.ImageFolder(filepath+'train', transform=data_transform['train']),  # 读取训练集'val'   : datasets.ImageFolder(filepath+'val', transform=data_transform['val'])  # 读取验证集}# 构造数据集
dataloader = {'train' : DataLoader(datasets['train'], batch_size=batch_size, shuffle=True),  # 构造训练集'val'   : DataLoader(datasets['val'], batch_size=batch_size, shuffle=True)   # 构造验证集}# --------------------------------------------- #
#（2）查看数据集信息
# --------------------------------------------- #
train_num = len(datasets['train'])  # 查看训练集数量
val_num = len(datasets['val'])  # 查看验证集数量# 查看分类类别及其索引  {0: 'cats', 1: 'dogs', 2: 'panda'}
class_names = dict((v,k) for k,v in datasets['train'].class_to_idx.items())
print(class_names)# 从训练集中取出一个batch，接收图片及其标签
train_imgs, train_labels = next(iter(dataloader['train']))
# 查看图像的标签及其shape  [32, 3, 224, 224]  [32]
print('img:', train_imgs.shape, 'labels:', train_labels.shape)

3.2 数据可视化

可视化训练集中的前12张图像。由于构造数据集时使用了一系列预处理方法，因此这里要将像素类型从 tensor 变成 numpy，调整图像的维度 [b,c,h,w]==>[b,h,w,c]，对图像的每个通道执行反标准化操作，恢复到0-1之间的随机分布。

标准化： $img = \frac{img-mean}{std}$ 反标准化： $img = img * std + mean$

代码如下：

# --------------------------------------------- #
#（3）数据可视化
# --------------------------------------------- #
# 从数据集中取出12张图及其标签
frames = train_imgs[:12]
frames_labels = train_labels[:12]# 将数据类型从tensor变回numpy
frames = frames.numpy()
# 维度调整 [b,c,h,w]==>[b,h,w,c]
frames = np.transpose(frames, [0,2,3,1])# 对图像做反标准化处理
mean = [0.485, 0.456, 0.406]  # 均值
std = [0.229, 0.224, 0.225]   # 标准化
# 图像的每个通道的特征图乘标准化加均值
frames = frames * std + mean# 将像素值限制在0-1之间
frames = np.clip(frames, 0, 1)# 绘制12张图像及其标签
plt.figure()  # 创建画板
for i in range(12):plt.subplot(3,4,i+1)plt.imshow(frames[i])plt.axis('off')  # 不显示坐标刻度plt.title(class_names[frames_labels[i].item()])  # 显示每张图片的标签
plt.tight_layout()  # 轻量化布局
plt.show()

查看训练集中的图像

3.3 模型加载，迁移学习

首先加载预训练权重 torch.load() 到内存中。由于预训练模型的分类数有1000个，即最后一个全连接层有 1000 个神经元，因此我们只用预训练权重的特征提取部分，不需要分类层部分。

遍历预训练权重文件，删除分类层 'head.weight', 'head.bias' 的权重。

这里不冻结预训练权重，所有权重参数都能通过反向传播更新。

代码如下：

# --------------------------------------------- #
#（4）模型加载，迁移学习
# --------------------------------------------- #
# 接收VIT模型，三分类
model = VIT(num_classes=3)# 加载预训练权重文件，文件中的分类层神经元个数是1k
pre_weights = torch.load(weightpath, map_location=device)# 删除权重文件中不需要的层，保留除了分类层以外的所有层的权重
del_keys = ['head.weight', 'head.bias']# 删除字典中的对应key
for k in del_keys:del pre_weights[k]# 将修改后的权重加载到模型上
# 当strict=True,要求预训练权重层数的键值与新构建的模型中的权重层数名称完全吻合
missing_keys, unexpected_keys = model.load_state_dict(pre_weights, strict=False)
print('miss:', len(missing_keys), 'unexpected:', len(unexpected_keys))# model.parameters() 代表网络的所有参数
for params in model.parameters():params.requires_grad = True  # 所有权重参与训练可以更新

3.4 模型训练

接下来进行网络训练，将所有需要计算的部分都搬运到 GPU 上，加快训练速度。

我这里使用每个epoch的验证集损失作为网络监控指标，如果损失小于规定值且一直在下降就保存当前 epoch 的权重。

还要注意的就是网络训练和测试的模式不一样，训练时 Dropout 层随机杀死神经元，BN 层取一个batch的均值和方差；验证时 Dropout 层不起作用，BN 层取整个训练集计算得到的均值和方差。通过 net.train() 和 net.eval() 来切换训练和验证模式。

代码如下：

# --------------------------------------------- #
#（5）网络编译
# --------------------------------------------- #
# 将模型搬运至GPU上
model.to(device)
# 定义交叉熵损失
loss_function = nn.CrossEntropyLoss()# 获取所有需要梯度更新的权重参数
params_optim = []
# 遍历网络的所有权重
for p in model.parameters():if p.requires_grad is True:  # 查看权重是否需要更新params_optim.append(p)   # 保存所有需要更新的权重print('训练参数：', len(params_optim))# 定义优化器，定义学习率，动量，正则化系数
optimizer = optim.SGD(params_optim, lr=0.001, momentum=0.9, weight_decay=3e-4)# --------------------------------------------- #
#（6）训练阶段
# --------------------------------------------- #
for epoch in range(epochs):print('='*30)  # 显示当前是第几个epoch# 将模型设置为训练模式model.train()# 记录一个epoch的训练集总损失total_loss = 0.0# 每个step训练一个batch，每次取出一个数据集及其标签for step, (images, labels) in enumerate(dataloader['train']):# 将数据集搬运到GPU上images, labels = images.to(device), labels.to(device)# 梯度清零，因为梯度是累加的optimizer.zero_grad()# 前向传播==>[b,3]logits = model(images)  # 得到每张图属于3个类别的分数#（1）损失计算# 计算每个step的预测值和真实值的交叉熵损失loss = loss_function(logits, labels)# 累加每个step的损失total_loss += loss#（2）反向传播# 梯度计算loss.backward()# 梯度更新optimizer.step()# 每50个epoch打印一次损失值if step % 50 == 0:print(f'step:{step}, train_loss:{loss}')# 计算一个epoch的训练集平均损失train_loss = total_loss / len(dataloader['train'])# --------------------------------------------- #
#（7）验证训练
# --------------------------------------------- #model.eval()  # 切换到验证模式total_val_loss = 0.0  # 记录一个epoch的验证集总损失total_val_correct = 0  # 记录一个epoch中验证集一共预测对了几个with torch.no_grad():  # 接下来不计算梯度# 每个step验证一个batchfor (images, labels) in dataloader['val']:# 将数据集搬运到GPU上images, labels = images.to(device), labels.to(device)# 前向传播[b,c,h,w]==>[b,3]logits = model(images)#（1）计算损失# 计算每个batch的预测值和真实值的交叉熵损失loss = loss_function(logits, labels)# 累加每个batch的损失，得到一个epoch的总损失total_val_loss += loss#（2）计算准确率# 找到预测值对应的最大索引，即该图片对应的类别pred = logits.argmax(dim=1)  # [b,3]==>[b]# 比较预测值和标签值，计算每个batch有多少预测对了val_correct = torch.eq(pred, labels).float().sum()# 累加每个batch的正确个数，计算整个epoch的正确个数total_val_correct += val_correct# 计算一个epoch的验证集的平均损失和平均准确率val_loss = total_val_loss / len(dataloader['val'])val_acc = total_val_correct / val_num# 打印每个epoch的训练集平均损失，验证集平均损失和平均准确率print('-'*30)print(f'train_loss:{train_loss}, val_loss:{val_loss}, val_acc:{val_acc}')# --------------------------------------------- #
#（8）保存权重
# --------------------------------------------- ## 保存最小损失值对应的权重文件if val_loss < best_loss:# 权重文件名称savename = savepath + f'epoch{epoch}_valacc{round(val_acc.item()*100)}%_' + 'VIT.pth'            # 保存该轮次的权重torch.save(model.state_dict(), savename)# 切换最小损失值best_loss = val_loss# 打印结果print(f'weights has been saved, best_loss has changed to {val_loss}')

训练过程如下：

==============================
step:0, train_loss:0.9088920950889587
step:50, train_loss:2.3867087364196777
step:100, train_loss:2.1412224769592285
------------------------------
train_loss:1.7520136833190918, val_loss:2.2571213245391846, val_acc:0.5276381969451904
==============================

训练过程中保存权重：

4. 预测阶段

接下来我们用训练好了的权重文件来预测图像的类别。同样先导入所有需要用到的工具包。

代码如下：

import torch
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from PIL import Image
from VisionTransformer_model import VIT
import numpy as np
import matplotlib.pyplot as plt# -------------------------------------------------- #
#（0）参数设置
# -------------------------------------------------- #
batch_size = 32  # 每次测试32张图# 测试集文件夹所在位置
file_path = 'D:/deeplearning/test/数据集/animal/test'
# 权重参数路径
weights_path = 'D:/deeplearning/imgnet/pytorchimgnet/save_weights/epoch5_valacc59%_VIT.pth'# 获取GPU设备
if torch.cuda.is_available():  # 如果有GPU就用，没有就用CPUdevice = torch.device('cuda:0')
else:device = torch.device('cpu')

4.1 构造数据集

这里测试集的预处理采用和验证集相同的预处理方法。这部分和上面相同，就不多做介绍。

# -------------------------------------------------- #
#（1）构造测试集
# -------------------------------------------------- #
# 定义测试集的数据预处理方法
data_transforms = transforms.Compose([transforms.Resize((224,224)),  # 将输入图像的size缩放至224*224transforms.ToTensor(),  # numpy边tensor，像素归一化，维度调整transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])  # 对每个通道标准化])# 加载测试集，并预处理
datasets = datasets.ImageFolder(file_path, transform=data_transforms)# 构造测试集
dataloader = DataLoader(datasets, batch_size=batch_size, shuffle=True)# 查看测试集一共有多少张图
test_num = len(datasets)# 获取测试集的分类类别及其索引  {0: 'cats', 1: 'dogs', 2: 'panda'}
class_names = dict((v,k) for k,v in datasets.class_to_idx.items())

4.2 多图像预测

取测试集的每个 Batch 的前12张图片，查看其预测结果。这里也需要做一次反标准化操作，和上面第三小节中相同。

在网络前向传播之前将模型设置为验证模式 model.eval()，只做前向传播的操作，不进行梯度更新操作 with torch.no_grad() 不计算梯度。

经过前向传播后，图像的shape变成 [b,3]，即图片预测属于3种类别的分数，然后经过softmax()求出图片属于每个类别的概率，通过torch.max()找出最大概率及其索引，得到图片属于哪个类别。

代码如下：

# -------------------------------------------------- #
#（2）绘图展示预测结果
# imgs:代表输入图像[b,c,h,w]，labels代表图像的真实标签[b]
# cls:代表每张图属的类别索引[b]，scores:代表每张图的类别概率[b]
# -------------------------------------------------- #
def im_show(imgs, labels, cls, scores):# 从数据集中取出12张图及其标签索引、概率frames = imgs[:12]true_labels = labels[:12]pred_labels = cls[:12]pred_scores = scores[:12]# 将数据类型从tensor变回numpyframes = frames.numpy()# 维度调整 [b,c,h,w]==>[b,h,w,c]frames = np.transpose(frames, [0,2,3,1])# 对图像做反标准化处理mean = [0.485, 0.456, 0.406]  # 均值std = [0.229, 0.224, 0.225]   # 标准化# 图像的每个通道的特征图乘标准化加均值frames = frames * std + mean# 将像素值限制在0-1之间frames = np.clip(frames, 0, 1)# 绘制12张图像及其标签plt.figure()  # 创建画板for i in range(12):plt.subplot(3,4,i+1)plt.imshow(frames[i])plt.axis('off')  # 不显示坐标刻度# 显示每张图片的真实标签、预测标签、预测概率plt.title('true:'+class_names[true_labels[i].item()] + '\n' +'pred:'+class_names[pred_labels[i].item()] + '\n' +'scores:'+str(round(pred_scores[i].item(), 3)))  plt.tight_layout()  # 轻量化布局plt.show()# -------------------------------------------------- #
#（3）图像预测
# -------------------------------------------------- #
# 模型构建
model = VIT(num_classes=3)
# 加载权重文件
model.load_state_dict(torch.load(weights_path, map_location=device))
# 将模型搬运到GPU上
model.to(device)
# 模型切换成测试模式，切换LN标准化和dropout的工作方式
model.eval()# 测试阶段不计算梯度
with torch.no_grad():# 每次测试一个batchfor step, (imgs, labels) in enumerate(dataloader):# 将数据集搬运到GPU上images, labels = imgs.to(device), labels.to(device)# 前向传播==>[b,3]logits = model(images)# 求出图像属于哪个类别索引[b,3]==>[b]pred_cls = logits.argmax(dim=1)# 计算图像属于每个类别的概率[b,3]==>[b,3]predicts = torch.softmax(logits, dim=1)# 获取最大预测类别的概率[b,3]==>[b]predicts_score, _ = predicts.max(dim=1)# 绘制预测结果im_show(imgs, labels, pred_cls, predicts_score)

查看预测结果：图像标题是真实类别、预测类别、预测概率值

【图像分类案例】(10) Vision Transformer 动物图像三分类，附Pytorch完整代码相关推荐

基于Vision Transformer的图像去雾算法研究与实现(附源码)
基于Vision Transformer的图像去雾算法研究与实现 0. 服务器性能简单监控 \LOG_USE_CPU_MEMORY\文件夹下的use_memory.py文件可以实时输出CPU使用率以及 ...
【机器视觉案例】(5) AI视觉，手势调节物体尺寸，附python完整代码
各位同学好,今天和大家分享一下如何使用opencv+mediapipe完成远程手势调节图片尺寸的案例.先放张图看效果.当拇指和食指竖起时,根据食指间的连线的长度自由缩放图片尺寸.图片的中点始终位于指尖 ...
【机器视觉案例】(5) AI视觉，远程手势控制虚拟计算器，附python完整代码
各位同学好,今天和大家分享一下如何使用MediaPipe+Opencv完成虚拟计算器,先放张图看效果.FPS值为29,食指和中指距离小于规定阈值则认为点击按键,为避免重复数字出现,规定每20帧可点击一 ...
【邢不行|量化小讲堂系列20-Python量化入门】10年400倍策略分享（附视频逐行代码讲解）
引言: 邢不行的系列帖子"量化小讲堂",通过实际案例教初学者使用python进行量化投资,了解行业研究方向,希望能对大家有帮助. [历史文章汇总]请点击此处 [必读文章]EOS期现 ...
Python 编程实现图像分辨率的修改（完整代码）
Python 编程实现图像分辨率的修改(完整代码) 在日常生活和工作中,我们经常需要对图片进行处理,其中一个常见的需求就是修改图像的分辨率.Python 语言提供了强大的图像处理库 PIL(Pytho ...
【神经网络】(3) 卷积神经网络(CNN)，案例：动物三分类，附python完整代码
各位同学好,今天和大家分享一下TensorFlow2.0深度学习中卷积神经网络的案例.现在有猫.狗.熊猫图片一千张,构建卷积神经网络实现图像的分类预测. 1. 数据加载将训练测试数据划分好后放在同一 ...
【图像分类案例】(1) ResNeXt 交通标志四分类，附Tensorflow完整代码
各位同学好,今天和大家分享一下如何使用 Tensorflow 构建 ResNeXt 神经网络模型,通过案例实战 ResNeXt 的训练以及预测过程.每个小节的末尾有网络.训练.预测的完整代码.想要数据 ...
【图像分类案例】(2) DenseNet 天气图片四分类（权重迁移学习），附Tensorflow完整代码
各位同学好,今天和大家分享一下使用 Tensorflow 构建 DenseNet 卷积神经网络模型,并使用预训练模型的权重,完成对四种天气图片的分类. 完整代码在我的 Gitee 中,有需要的自取: ...
【神经网络】(2) 网络优化，案例：服装图像分类，附python完整代码
各位同学好,今天和大家分享以下TensorFlow2.0深度学习中对神经网络的优化方法,包括动量.学习率.dropout.交叉验证.正则化.本节使用数学公式对网络进行优化,增加网络的灵活性. 以下代码 ...

【图像分类案例】(10) Vision Transformer 动物图像三分类，附Pytorch完整代码

1. 引言

2. 模型构建

2.1 Patch Embedding

2.2 类别标签和位置编码

2.3 多头自注意力模块

2.4 MLP 多层感知器

2.5 特征提取模块

2.6 主干网络

3. 训练阶段

3.1 构造数据集

3.2 数据可视化

3.3 模型加载，迁移学习

3.4 模型训练

4. 预测阶段

4.1 构造数据集

4.2 多图像预测

【图像分类案例】(10) Vision Transformer 动物图像三分类，附Pytorch完整代码相关推荐

最新文章

热门文章