pytorch官方demo介绍(使用CIFAR10数据集)
1.模型的构建
from torch import nn as nn
import torch.nn.functional as F
class model_LeNet(nn.Module):
def __init__(self):
super(model_LeNet, self).__init__()
self.conv1 = nn.Conv2d(3,16,5) #inputchannel , outputchannel, kernel_size
self.pool1 = nn.MaxPool2d(2,2) #
self.conv2 = nn.Conv2d(16,32,5) #inputchannel , outputchannel, kernel_size
self.pool2 = nn.MaxPool2d(2,2)
self.fc1 = nn.Linear(32*5*5,120)
self.fc2 = nn.Linear(120,84)
self.fc3 = nn.Linear(84,10)
def forward(self,x):
x = F.relu(self.conv1(x)) #input(3,32,32) output(16,28,28)
x = self.pool1(x) #input(16,28,28) output(16,14,14)
x = F.relu(self.conv2(x)) #input(16,14,14) output(32,10,10)
x = self.pool2(x) #input(32,10,10) output(32,5,5)
x = x.view(-1,32*5*5) #第一个参数batch是自动计算的,第二个维度是把他展成一维的,跟flaten差不多,只是在转换类型上不一样
#output(batch,32*5*5)的一维向量
x = F.relu(self.fc1(x)) #output(120)
x = F.relu(self.fc2(x)) #output(84)
x = self.fc3(x) #output(10)
return x
#测试上述构建的模型
import torch
#随意生成一个32张图片channel为3通道且高和宽均为32的图片集
input1 = torch.rand([32,3,32,32])
#创建模型
lenet = model_LeNet()
#打印出模型的具体内容
print(lenet)
#把数据输入进去得到输出
output = lenet(input1)
#展示出输出的内容
print(output)
首先,model_LeNet
类继承自 nn.Module
,这是PyTorch中构建神经网络的基本类。在 __init__
方法中:
- 定义了两个卷积层(
conv1
和conv2
),分别将输入的3通道图像转换为16和32通道特征图,使用5x5大小的卷积核。 - 两个最大池化层(
pool1
和pool2
)用于减小特征图的空间尺寸,采用2x2窗口大小且步长为2的Max Pooling操作。 - 定义了三个全连接层(
fc1
,fc2
,fc3
),前两个之后接ReLU激活函数,最后一层直接输出分类结果,类别数为10(对应CIFAR-10数据集中的10个类别)。
在 forward
方法中,描述了前向传播过程:
- 对输入数据进行两次卷积及激活操作,每次后跟一次最大池化操作,逐步提取特征并减小空间维度。
- 将经过第二次池化后的特征图展平为一维向量,输入到全连接层。
- 经过三层全连接层处理,最后得到一个形状为
(batch_size, 10)
的输出张量,表示每个样本属于10个类别中的每一个类别的预测得分。
测试部分:
- 随机生成了一组形状为
[32, 3, 32, 32]
的四维张量,代表一个批量大小为32、每张图片具有3个通道、高和宽均为32像素的数据集。 - 创建一个
model_LeNet
实例,并打印模型结构。 - 将随机生成的数据送入模型进行前向传播计算,得到输出张量。
- 打印模型的输出内容,即对这个随机生成数据集的预测结果。
2.训练数据集
import torch
import torchvision.transforms as transforms
import torchvision.datasets
from torch import nn
from torch.utils.data import DataLoader
from LeNet_pytroch_office.model_LeNet import model_LeNet
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 训练数据集
trainset = torchvision.datasets.CIFAR10("./dataset",train=True,
transform=transform,download=True)
trainloader = torch.utils.data.DataLoader(trainset,batch_size=36,shuffle=True,num_workers=0)
# 测试数据集
testset = torchvision.datasets.CIFAR10("./dataset",train=False,
transform=transform,download=True)
testloader = torch.utils.data.DataLoader(testset,batch_size=10000,shuffle=False,num_workers=0)
test_data_iter = iter(testloader)
test_image,test_lable = next(test_data_iter)
classes = ('plane','car','bird','cat','deer','dog','frog','horse','ship','truck')
#创建模型实例
net = model_LeNet()
#定义损失函数
loss_function = nn.CrossEntropyLoss()
#定义优化器
optimizer = torch.optim.Adam(net.parameters(),lr=0.001)
print("start:")
for epoch in range(5): #数据集循环的次数,也成为训练轮数
running_loss = 0.0
for step,data in enumerate(trainloader,start=0):
imgs,labels = data
#消除梯度,设梯度为0
optimizer.zero_grad()
# 进行处理图片
outputs = net(imgs)
#计算损失函数
loss = loss_function(outputs,labels)
#进行反向传播
loss.backward()
#使用优化器进行参数更新
optimizer.step()
#打印状态
running_loss += loss.item()
if step%500 ==499:
with torch.no_grad():
outputs = net(test_image)
predict_y = torch.max(outputs,dim=1)[1]
accuracy = (predict_y == test_lable).sum().item() / test_lable.size(0)
print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f'%(epoch+1,step+1,running_loss/500,accuracy))
running_loss = 0.0
print("finished Training")
save_path = './Lenet.pth'
torch.save(net.state_dict(),save_path)
- 首先导入所需库并定义图像预处理变换,包括将图像转换为Tensor以及进行归一化操作。
- 使用torchvision.datasets.CIFAR10加载CIFAR-10数据集,分别获取训练集和测试集。设置相应的数据预处理方式(transform)并在本地下载数据集(如果尚未存在)。
- 创建DataLoader以加载训练集和测试集数据,并设置合适的批大小、是否打乱数据(仅训练集需要)以及工作进程数。
- 初始化一个迭代器来访问测试集中的第一个批次数据。
- 定义类别名称列表(CIFAR-10的10个类别)。
- 创建一个
model_LeNet
实例作为网络模型,并定义损失函数(交叉熵损失)和优化器(Adam,学习率为0.001)。 - 开始训练过程,循环遍历5个epoch:
- 在每个epoch中,计算累积损失(running_loss),并对每个训练批次执行以下操作:
- 将梯度置零,准备进行反向传播。
- 将当前批次的数据输入网络并得到预测结果(outputs)。
- 计算当前批次的损失值。
- 进行反向传播求解梯度。
- 使用优化器更新网络参数。
- 每隔500个训练步骤,在无梯度计算环境下评估模型在测试集上的表现:
- 将测试集的第一个批次数据输入网络,获取预测类别。
- 计算准确率。
- 打印当前epoch、训练步数、平均训练损失和测试准确率。
- 在每个epoch中,计算累积损失(running_loss),并对每个训练批次执行以下操作:
- 训练完成后,保存模型权重至’./Lenet.pth’文件。
这段代码的主要功能是在CIFAR-10数据集上训练一个自定义的LeNet模型,并在训练过程中对测试集进行验证,最后保存训练好的模型权重。
3.进行测试(构建数据对数据进行预测)
图片是我随意在网上下载的,我是下载了十张图片,十张图片对应十个类别,我想看看十张图片能有多少张图片预测正确
我下载的图片放在了pycharm的同级目录下,如图
import torch
import torchvision.transforms as transforms
from PIL import Image
from model_LeNet import model_LeNet
# 定义一个图像预处理流水线 transforms,它包含三个转换步骤:
# 使用 transforms.Resize((32,32)) 将所有图像调整为32x32大小。
# 使用 transforms.ToTensor() 将 PIL 图像转换为 PyTorch 张量。
# 使用 transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5)) 对输入图像数据进行归一化,减去每个通道的均值(0.5)并除以标准差(0.5),
# 使得输入数据满足模型训练时的要求。
transforms = transforms.Compose(
[
transforms.Resize((32,32)),
transforms.ToTensor(),
transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
]
)
# 定义十种类型
classes = ('plane','car','bird','cat','deer','dog','frog','horse','ship','truck')
# 初始化一个LeNet模型实例 net = model_LeNet(),然后加载预训练好的权重参数 net.load_state_dict(torch.load('Lenet.pth'))。
net = model_LeNet()
net.load_state_dict(torch.load('Lenet.pth'))
for i in range(10):
# 打开对应的图像文件,文件名为(i + 1).jpg。
path = str((i+1))+'.jpg'
image = Image.open(path)
# 使用之前定义的 transforms 对图像进行预处理。
image = transforms(image)
# 通过 torch.unsqueeze(image, dim=0) 添加一个额外维度,
# 以便输入到网络中(因为神经网络通常需要输入批次数据,即使只有一个样本也需要在第一个维度上具有批量大小)
image = torch.unsqueeze(image, dim=0)
# 无梯度计算的环境下(即 with torch.no_grad(): 内部),将预处理后的图像输入到网络中得到输出结果 outputs。
with torch.no_grad():
outputs = net(image)
# 使用 torch.max(outputs, dim=1)[1].data.numpy() 获取预测类别对应的索引(最大概率对应的类别)。
# predict = torch.max(outputs, dim=1)[1].data.numpy()
#使用softmax函数进行处理
predict = torch.softmax(outputs,dim=1)
#进行了softmax函数处理,dim=1表示在特征数上进行softmax处理,把各个特征的值转换成[0,1]的概率分布。
print(predict)
#根据预测的索引值从 classes 列表中获取并打印对应的类别名称。
predict = torch.max(outputs, dim=1)[1].data.numpy()
print(classes[int(predict)])
- 首先定义了一个图像预处理流水线
transforms
,它将图像调整为32x32大小、转换为PyTorch张量格式,并对其进行归一化处理。 - 定义了10个类别标签列表
classes
,这些类别与模型要预测的目标相对应。 - 初始化一个LeNet模型实例
net = model_LeNet()
并加载预训练权重net.load_state_dict(torch.load('Lenet.pth'))
。 - 进行循环,针对从0到9的每个索引:
- 根据索引生成文件路径并打开对应的图像文件。
- 使用之前定义的
transforms
对图像进行预处理。 - 通过
torch.unsqueeze(image, dim=0)
将处理后的图像添加一个维度以适应网络输入要求(批次维度)。 - 在无梯度计算环境下运行模型,获取输出结果
outputs
。 - 修改了原本的预测逻辑:原代码是直接获取最大概率对应的类别索引;现在使用
torch.softmax(outputs, dim=1)
计算每个类别的softmax概率分布。
- 打印出经过softmax函数处理后的预测概率分布
predict
。
这段代码的核心是对十个不同的.jpg格式图片进行分类预测,利用已训练好的LeNet模型,将其转换为模型所需格式并进行前向传播,最终输出预测类别。