train.py : 训练代码
import torch.optim
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
'''
torchvision.datasets: 数据集模块,里面收集了若干个数据集类型,CIFAR10 就是其中一个
root:数据集所在目录的根目录 如果download设置为True。“cifar-10-batches-py '”存在,则将被保存至该目录
train :如果为True,则从训练集加载数据集,否则从测试集加载。
transform::(bool,可选)一个接受PIL图像的函数/变换 并返回转换后的版本。 如”transforms“,“RandomCrop”
torchvision.transforms.ToTensor():
将PIL(Python Imaging Library)或者 numpy.ndarray 类型的图像数据转换为 PyTorch 中的 tensor 类型。
神经网络的输入数据通常是 tensor 类型,此方法还会进行一些归一化操作。
具体操作:
将像素点的范围从 [0, 255] 归一化到 [0, 1]。
将像素点的数据类型从 uint8 转换为 float32。
对于彩色图像,将通道维度从 (H, W, C) 转换为 (C, H, W),其中 C 表示通道数,H 表示高度,W 表示宽度。
download:如果为True时,表示如果指定路径中不存在该数据集,则自动下载并存放在该路径中。
'''
# 准备数据集
# 训练数据集
train_data = torchvision.datasets.CIFAR10(root="../data", train=True, transform=torchvision.transforms.ToTensor(),
download=True)
# 测试数据集
test_data = torchvision.datasets.CIFAR10(root="../data", train=False, transform=torchvision.transforms.ToTensor(),
download=True)
# 数据集长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集长度:{}".format(train_data_size))
print("测试数据集长度:{}".format(test_data_size))
'''
DataLoader:
创建数据加载器对象、对训练数据集进行批量加载,实现数据的随机打乱、并行加载等操作。每次迭代训练时,
我们可以通过 train_dataLoader 对象来获取一个 batch 的数据,并将其输入到模型中进行训练,从而逐步优化模型的参数。
train_data:表示要加载的训练数据集对象。
batch_size=64:表示每个 batch 的大小,默认为 1,这里设置为 64。表示一次性加载64张图片
'''
# 利用 DataLoader 来加载数据集
train_dataLoader = DataLoader(train_data, batch_size=64)
test_dataLoader = DataLoader(test_data, batch_size=64)
'''
nn.Module :PyTorch 中的一个核心概念,是神经网络模型的基本构建块。
所有的神经网络模型都需要继承自 nn.Module 类,并重写其中的 __init__() 和 forward() 函数。
nn.Sequential:可以用来构建神经网络模型,其作用是将一系列的神经网络层按照顺序串联在一起,构成一个完整的神经网络模型。
nn.Conv2d(3, 32, 5, 1, 2):二维卷积层、将数据输入进行卷积操作。其参数为:
第一个参数(3): in_channels:输入数据的通道数,例如 RGB 图像的通道数为 3,灰度图像的通道数为 1。
第二个参数(32): out_channels:卷积核的数量,也即输出特征图的通道数。
第三个参数(5): kernel_size:卷积核的大小,可以是一个整数,表示正方形卷积核的边长,也可以是一个二元组,表示长宽不同的矩形卷积核的长宽。
第四个参数(1): stride:卷积核的步长,表示每次卷积操作时卷积核在输入特征图上移动的距离。
第五个参数(2): padding:卷积核的填充大小,用于控制输出特征图的大小。如果设置为 0,则表示不进行填充;如果设置为 k,则表示在输入特征图的边缘填充 k 个像素,使得卷积核可以顺利进行卷积操作。
nn.MaxPool2d(2): 表示按照最大值进行池化、参数 kernel_size = 2、表示池化核大小(2X2),即将原来图像池化为原来一半
nn.Flatten():将多维数据变成一维、让全连接层能够使用
nn.Linear(64*4*4, 64):经过上面将数据变为一维后,一维数据长度为64*4*4(经过计算得出,不同图片,不同神经网络可能会不同),
此层将64*4*4长度的数据转换为长度为64的数据
forward(self, x):必须重写的方法,作用是前向传播,通过此方法获取模型输出结果,具体来说:
x 是模型的输入,代表一个 batch 的图片数据。在 forward() 方法中,首先将输入 x 传入 self.model 中,经过一系列的卷积、池化和线性变换操作后,
最终得到输出结果 x。在这个示例中,模型的输出 x 是一个大小为 (batch_size, 10) 的张量,代表每个输入样本对应的类别概率。
'''
# 创建网络模型(有两种方式,第一种是什么都没有时,就从0创建,第二种是已经有模型保存了,就从保存的模型中提取,此处是第一种)
class Zou(nn.Module):
def __init__(self):
super(Zou, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4, 64), # 将长度从 64*4*4 变为 64
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
zou = Zou()
# 第二种模型加载方式,前提一定要定义一个相同的模型结构用来接收
zou.load_state_dict(torch.load("zou_9.pth"))
# 判断是否能使用GPU进行训练
if torch.cuda.is_available():
zou = zou.cuda()
'''
nn.CrossEntropyLoss():损失函数,用于分类问题,将预测结果与真实标签进行比较,并计算两者的交叉熵损失,函数中已经包含了 Softmax 操作的计算过程,具体来说:
nn.CrossEntropyLoss() 的输入是一个大小为 (batch_size, num_classes) 的张量 x 和一个大小为 (batch_size,) 的张量 y,
其中 x 表示模型对于当前 batch 的预测结果,y 表示当前 batch 中样本的真实标签。在计算损失时,
nn.CrossEntropyLoss() 首先对 x 进行 nn.LogSoftmax() 操作,得到一个新的张量,然后将这个新的张量和真实标签 y 传入 nn.NLLLoss() 函数中,
计算交叉熵损失。
nn.LogSoftmax():将一个向量(一张图相当于二维向量),经过Softmax映射为概率,如向量[0.5, -1.2, 3.8]可以被映射为:(0.018, 0.004, 0.978)
nn.NLLLoss():计算模型输出与真实标签之间的距离,即我们知道,这个函数的输入会是 outputs 和 targets ,我们又知道,outputs是一个二维向量,
每一个向量里面有10数,这10个数对应是不同动物的概率,而targets只是一维向量,里面的每一个数表示每一张图片的正确结果,所以这函数的作用是,
假如:我们知道正确结果是 标签5,即狗这个动物,那么这个就可以类似为一个有10数的一维向量,此向量里面的值除了第五个数是1外,其余皆是0,(1表示概率为100%),
那么我们取outputs和此一维向量中的第五位数进行相减,就表示损失值
'''
# 损失函数
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
'''
torch.optim.SGD:随机梯度下降、每次迭代中只考虑一个batch数据,根据当前数据的损失函数梯度来更新模型的参数
zou.parameters():返回的是神经网络模型zou中的所有可训练参数,即需要进行梯度更新的参数,这些参数包括nn.Module对象的所有参数。
torch.optim.SGD()函数的第一个参数需要传入要进行参数更新的可训练参数。
'''
# 优化器
learning_rate = 0.01 # 梯度下降的步长
optimizer = torch.optim.SGD(zou.parameters(), lr=learning_rate)
# 设置训练网络的参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10
# 添加 tensorboard ,即训练结果图像化
writer = SummaryWriter("../logs_train")
'''
它首先从 train_dataLoader 中迭代取出一个 batch 的数据,然后将输入数据 imgs 和标签数据 targets 送入模型 zou 中得到模型的输出结果 outputs。
接着,将模型输出结果 outputs 与标签数据 targets 一起输入到损失函数 loss_fn 中,计算出当前 batch 的训练损失 loss。
然后,通过优化器 optimizer 对模型进行优化,即将损失 loss 反向传播回模型中,求得模型中每个参数的梯度,
然后使用优化器对每个参数进行更新,最终使得损失 loss 越来越小,模型的表现也越来越好。
targets: 表示真实结果,如我们一次性训练64张图片,那么targets的值为:
tensor([6, 9, 9, 4, 1, 1, 2, 7, 8, 3, 4, 7, 7, 2, 9, 9, 9, 3, 2, 6, 4, 3, 6, 6,
2, 6, 3, 5, 4, 0, 0, 9, 1, 3, 4, 0, 3, 7, 3, 3, 5, 2, 2, 7, 1, 1, 1, 2,
2, 0, 9, 5, 7, 9, 2, 2, 5, 2, 4, 3, 1, 1, 8, 2])
每一个值代表一张图片属于哪一个类型,如6,代表青蛙等等(标签代表什么动物是事先定义好的、我们刚开始的模型是不知道这个标签代表什么,
只是说如果模型预测的结果刚好与标签相同,那么就告诉模型,你正确了,然后模型就记住了,如此循环往复慢慢调教)
outputs:
直接看输出:
tensor([
[-3.1442e+00, -4.2692e+00, 1.5371e+00, 3.6352e+00, 1.5766e+00,3.7046e+00, 5.7747e+00, 1.8751e+00, -6.9454e+00, -6.6534e+00],
[ 9.3670e-01, 3.2580e+00, 2.1656e+00, -1.3350e+00, -2.7406e+00,-9.9861e-01, -8.0962e+00, 4.5565e+00, -2.3388e+00, 6.2317e+00],
,,,
[-2.6294e-01, 7.9632e-01, 2.6810e+00, 2.5335e-01, -1.4035e+00, 7.0486e-01, -1.2330e+00, 6.4206e-01, -2.2874e+00, -1.1087e+00]],
device='cuda:0', grad_fn=)
每一行是一个一维数组,有10个数,每个数表示一个类型对应的预测值,哪个预测值最大表示是对应动物的可能性越大(数组下标加一就代表对应动物的预测值,如 1.5766e+00:表示狗的预测值)
一共有64行,代表对64张图片的预测,每一行取一个最大值就表示预测最终结果,
但是这个行为会发生在 accuracy = (outputs.argmax(1) == targets).sum() 函数里面(预测时使用)
即,先横向取数组中最大值,然后通过最大值所在的下标加1与对应的targets值比较,判断预测是否正确,然后将所有正确的加起来
'''
for i in range(epoch):
print("----------第 {} 轮训练开始".format(i))
# 训练步骤开始
zou.train() # 固定写法,用于将模型 zou 切换到训练模式,会对特定网络或参数有加速效果
for data in train_dataLoader:
imgs, targets = data
# 判断是否能使用GPU进行加速
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = zou(imgs) # 将图像放入模型中进行预测
loss = loss_fn(outputs, targets) # 训练结果 outputs 与 真实结果 targets 进行比较并返回 损失loss。
# 优化器优化模型
optimizer.zero_grad() # 梯度清零、在PyTorch中,每一次计算梯度都会累加到之前的梯度上,因此在每次反向传播前,需要手动清除上一次计算的梯度,避免对本次梯度计算的影响。
loss.backward() # 是 PyTorch 中计算张量梯度的函数、根据链式法则自动计算梯度,即将当前的梯度值传递给前面的层,通过链式法则计算出每一层的梯度。
optimizer.step() # 使用优化器的 step() 函数根据梯度来更新参数,从而实现梯度下降优化。
total_train_step += 1 # 训练次数加1
if total_train_step % 100 == 0:
print("训练次数:{},loss:{}".format(total_train_step, loss.item()))
writer.add_scalar("train_loss",loss.item(), total_train_step)
# 测试步骤开始
zou.eval() # 固定写法,表示开始测试,针对某些内容有效果
total_test_loss = 0 # 此轮测试总偏差值
total_accuracy = 0
'''
是一个上下文管理器,用于指定在该上下文中,PyTorch 不会记录梯度信息,也就是不进行反向传播。在该上下文中进行前向传播,可以节省显存,并且速度较快。通常用于测试阶段的推断(inference)。
'''
with torch.no_grad():
for data in test_dataLoader:
imgs, targets = data
# 判断是否能使用GPU进行加速
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = zou(imgs)
loss = loss_fn(outputs, targets)
total_test_loss += loss.item()
'''
outputs.argmax(1):因为输出是一个长度为10的一维向量,此方法参数为1表示按行取最大值然后与targets进行比较(即)
'''
accuracy = (outputs.argmax(1) == targets).sum() # 将得到的数组
total_accuracy += accuracy
print("整体测试集上的Loss:{}".format(total_test_loss))
print("整体测试集上的正确率:{}".format(total_accuracy / test_data_size))
writer.add_scalar("test_loss", total_test_loss, total_test_step) # total_test_step:测试次数
writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
total_test_step += 1
# 保存模型(每一轮保存一次)(只保存模型参数,使用字典模式)
torch.save(zou.state_dict(), "zou_{}.pth".format(i))
print("模型已保存")
writer.close()
测试代码
test.py
import torch
import torchvision.transforms
from PIL import Image
from torch import nn
image_path = "img.png"
image = Image.open(image_path) # 函数会返回一个'PIL.图片.图片PIL.Image.Image对象,表示打开的图片
image = image.convert('RGB') # 保证图片是三通道的
'''
torchvision.transforms.Compose:是一个将多个图像预处理步骤组合在一起的类,用于构建预处理管道。它将多个预处理操作封装在一起,以便可以在训练和测试期间一次性应用它们
torchvision.transforms.Resize((32, 32)):将输入图像调整为指定大小,即32X32
torchvision.transforms.ToTensor():将图像的数据形式转换为张量形式,
如果一张图片大小为32X32,那么转换为张量形式则是(3*32*32,)即从一张彩色图片转化为一个三维数组(3*256*256)其中第一个维度对应三个通道,后面两个维度分别对应图像的高度和宽度。
transform(image) :将输入的图片进行如上操作
'''
# 处理图片大小
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
torchvision.transforms.ToTensor()])
image = transform(image)
print(image.shape)
# 创建网络模型
class Zou(nn.Module):
def __init__(self):
super(Zou, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4, 64), # 将长度从 64*4*4 变为 64
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
zou = Zou()
zou.load_state_dict(torch.load("zou_9.pth")) # 加载已经有的模型
print(zou)
'''
torch.reshape(image, (1, 3, 32, 32)):
这段代码将张量image的形状从原来的(3, 32, 32)变成了(1, 3, 32, 32)。其中:
1是新的第一维,表示这是一个batch size为1的数据(即一次性只处理一张图片);
3表示图像的通道数,即RGB三通道;
32和32表示图像的高和宽。
'''
image = torch.reshape(image, (1, 3, 32, 32))
zou.eval() # 固定步骤表示开始测试
with torch.no_grad():
output = zou(image)
print(output)
print(output.argmax(1))