Pytorch之CIFAR10分类卷积神经网络
- 个人主页:风间琉璃
- 版权: 本文由【风间琉璃】原创、在CSDN首发、需要转载请联系博主
- 如果文章对你有帮助、欢迎关注、点赞、收藏(一键三连)和订阅专栏哦
前言
在计算机视觉领域，图像分类是一个重要的问题。CIFAR-10是一个常用的图像分类数据集，包含10个类别的60000张32x32彩色图像。在本篇博客中，我们将使用PyTorch构建一个卷积神经网络（Convolutional Neural Network，CNN）来对CIFAR-10数据集中的图像进行分类。
一、CIFAR-10数据集
CIFAR-10数据集包含10个类别，每个类别有6000张32x32大小的彩色图像。这些类别分别是飞机、汽车、鸟类、猫、鹿、狗、青蛙、马、船和卡车。数据集被分为训练集和测试集，每个集合各有50000张图像和10000张图像。
二、实现卷积神经网络
我们将构建一个简单的卷积神经网络来进行图像分类。以下是网络的架构：
- 卷积层（Convolutional Layer）：使用5x5的卷积核进行特征提取。
- 池化层（Pooling Layer）：使用2x2的最大池化操作来减小特征图的大小。
- 全连接层（Fully Connected Layer）：将池化后的特征图展平并连接到一个全连接层，即上图中绿色的线条, 用于分类处理。
- 输出层：得到每个类别的概率分布。
根据神经网络结构图，搭建网络模型时，需要计算的是卷积层中的stride和padding。根据官方提供的计算公式，如下：
以第一层卷积层为例子，in_channels = 3，out_channels=32; Hin=32，Win=32；Hout=32，Wout=32; kernel_size=5。
注意，当stride/padding/dilation为int时，其sride[0]=stride[1]=int，padding[0]=padding[1]=int，dilation[0]=dilation[1]=int。
将已知量带入公式，可计算得到stride=1，padding=2。
以下是网络的PyTorch代码实现：
from torch import nn
import torch
# CIFAR10网络模型
class Cifarnet(nn.Module):
def __init__(self):
super(Cifarnet, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1,
padding=2), # stride=1,padding=2卷积后尺寸不变
nn.MaxPool2d(2),
nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1,
padding=2), # stride=1,padding=2卷积后尺寸不变
nn.MaxPool2d(2),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1,
padding=2), # stride=1,padding=2卷积后尺寸不变
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
if __name__ == '__main__':
net = Cifarnet()
input = torch.ones((64, 3, 32, 32))
output = net(input)
# torch.Size([64, 10]) 每一张图片的预测概率为一个10大小的数组
print(output.shape)
三、数据集的加载与训练
(1)CPU训练
在构建网络之后，我们需要加载CIFAR-10数据集并进行训练。我们将使用PyTorch的数据加载工具来处理数据集。
# 准备数据集
train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10(root='./dataset', 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 = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)接下来，我们将定义损失函数和优化器，并使用训练数据对网络进行训练。
# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 优化器
learning_rate = 0.01
optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate)
# 记录训练网络的一些参数
# 记录训练次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 记录训练的轮数
epoch = 25
# 使用tensorboard可视化数据
writer = SummaryWriter(log_dir='./logs')
for i in range(epoch):
print("-----第{}轮训练开始-----".format(i+1))
# 开始训练
net.train() # 仅对特殊的网络层右作用
for data in train_dataloader:
imgs, targets = data
outputs = net(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step += 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)
(2)GPU训练
在使用GPU训练时，只需要将上面的代码作一些微调就可以了。 将cpu转换为gpu上，需要改三个地方：网络模型、损失函数、训练中加载的数据集和测试中加载的数据集
# 定义训练的设备
device = torch.device("cuda")
# 创建网络模型
net = Cifarnet()
net = net.to(device)
# 训练集和测试集一样的操作
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
四、测试与评估
# 开始测试
net.eval()
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
outputs = net(imgs)
loss = loss_fn(outputs, targets)
total_test_loss += loss.item()
# argmax找出每一张预测概率最大的下标
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)
writer.add_scalar("test_accurary", total_accuracy/test_data_size, total_test_step)
total_test_step += 1以下是使用CPU训练了5个epoch，所以数据并不是很理想，电脑配置高的建议100以上训练。
测试集正确率
测试集损失函数：
训练集损失函数：
使用Google colab gpu进行训练200个epoch：
训练200次正确率也不是很高，可能是由于没有调参的问题。
模型的保存，为后续实现分类作准备。
# 保存模型
torch.save(net, "cifdarnet.pth")
print("模型已保存")五、实现分类
利用刚刚训练完成的网络模型，进行图像的分类。
既然是为了实现图像分类，那么我们就要先获得图片对应的标签(分类类别)，在PyCharm中可以在数据集预处理的时候通过添加断点Debug的模式获取。
# ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
class_id = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']注意：如果使用GPU训练的模型，而使用CPU去加载测试，需要做如下处理：
将gpu改为cpu时，遇到一个报错：
RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location='cpu' to map your storages to the CPU.
model = torch.load("cifdarnet_gpu.pth", map_location='cpu')(1)以图片的形式测试
获取到类别后，然后就是将读取的图片进行预处理，其大小和信息要符合我们前面定义的网络模型。
# image_path = 'D:\Data\Learn_Pytorch\cifar10_train\dog.png'
image_path = 'D:\Data\Learn_Pytorch\cifar10_train\plane.png'
image_raw = Image.open(image_path)
image = image_raw.convert('RGB') # 将ARGB-->RGB
print(image) # [3,32,32]
# 网络的输入为32X32
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
torchvision.transforms.ToTensor()])
image = transform(image)
print(image.shape)预处理主要的操作有颜色空间类型转换，要符合RGB颜色空间，其次图片的缩放，图片的尺寸要符合32x32，并且图片的类型要是tensor类型。
然后就可以加载前面保存的网络模型进行图像的分类：
model = torch.load("cifdarnet.pth")
print(model)
# [3,32,32] --> [1,3,32,32]
image = torch.reshape(image,[1, 3, 32, 32])
model.eval()
with torch.no_grad():
output = model(image)
idx = output.argmax(1)
reslut = class_id[idx]
# 可视化
draw = ImageDraw.Draw(image_raw)
draw.text((10, 10), reslut, fill=(255, 0, 0))
image_raw.show()
print(output)
飞机预测结果：
小狗预测结果：
突然发现这么低的正确率，这个网络还是能识别处理的。
(2)以视频流的形式测试(OpenCV)
首先要了解各种不同加载图片的通道顺序
所以对摄像头传过来的视频流主要是作通道顺序转换和图片缩放，以及图片转换为tensor(float)。
# 将opencv读取的图片进行预处理后，送网络中
# opencv 读取的图像s缩放为32X32的大小
resized_img = cv2.resize(image_raw, (32, 32), i
# bgr-->rgb
image_rgb = cv2.cvtColor(resized_img, cv2.COLOR
# 转torch.Tensor
image = torch.from_numpy(image_rgb)
# 输入的图片是整数，要和网络参数(浮点数)保持一致
image = image.float()
# print(image.shape)
# [3,32,32] --> [1,3,32,32] 提升一张照片的维度
image = torch.reshape(image, [1, 3, 32, 32]) 汽车预测：
六、源码
(1)训练
网络模型：
from torch import nn
import torch
# CIFAR10网络模型
class Cifarnet(nn.Module):
def __init__(self):
super(Cifarnet, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1,
padding=2), # stride=1,padding=2卷积后尺寸不变
nn.MaxPool2d(2),
nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1,
padding=2), # stride=1,padding=2卷积后尺寸不变
nn.MaxPool2d(2),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1,
padding=2), # stride=1,padding=2卷积后尺寸不变
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
if __name__ == '__main__':
net = Cifarnet()
input = torch.ones((64, 3, 32, 32))
output = net(input)
# torch.Size([64, 10]) 每一张图片的预测概率为一个10大小的数组
print(output.shape)
网络训练和保存：
import torch
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from model import *
# 准备数据集
train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10(root='./dataset', 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 = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
# 创建网络模型
net = Cifarnet()
# print(net)
# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 优化器
learning_rate = 0.01
optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate)
# 记录训练网络的一些参数
# 记录训练次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 记录训练的轮数
epoch = 5
# 使用tensorboard可视化数据
writer = SummaryWriter(log_dir='./logs')
for i in range(epoch):
print("-----第{}轮训练开始-----".format(i+1))
# 开始训练
net.train() # 仅对特殊的网络层右作用
for data in train_dataloader:
imgs, targets = data
outputs = net(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step += 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)
# 开始测试
net.eval()
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
outputs = net(imgs)
loss = loss_fn(outputs, targets)
total_test_loss += loss.item()
# argmax找出每一张预测概率最大的下标
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)
writer.add_scalar("test_accurary", total_accuracy/test_data_size, total_test_step)
total_test_step += 1
# 保存模型
torch.save(net, "cifdarnet.pth")
print("模型已保存")
writer.close()(2)分类
图片实现(PIL)
import torch
import torchvision
from PIL import Image, ImageDraw
# ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
class_id = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
image_path = 'D:\Data\Learn_Pytorch\cifar10_train\dog.png'
#image_path = 'D:\Data\Learn_Pytorch\cifar10_train\plane.png'
image_raw = Image.open(image_path)
image = image_raw.convert('RGB') # 将ARGB-->RGB
print(image) # [3,32,32]
# 网络的输入为32X32
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
torchvision.transforms.ToTensor()])
image = transform(image)
print(image.shape)
model = torch.load("cifdarnet.pth")
print(model)
# [3,32,32] --> [1,3,32,32]
image = torch.reshape(image,[1, 3, 32, 32])
model.eval()
with torch.no_grad():
output = model(image)
idx = output.argmax(1)
reslut = class_id[idx]
# 可视化
draw = ImageDraw.Draw(image_raw)
draw.text((10, 10), reslut, fill=(255, 0, 0))
image_raw.show()
print(output)
摄像头实现：
import torch
import cv2
# ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
class_id = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
# 加载分类模型
model = torch.load("cifdarnet.pth")
# print(model)
# 捕获摄像头图像
cap = cv2.VideoCapture(1)
# 检查摄像头是否成功打开
if not cap.isOpened():
print("无法打开摄像头")
exit()
# 循环读取帧
while True:
# 逐帧捕获
ret, image_raw = cap.read()
# 检查帧是否读取成功
if not ret:
print("无法获取帧")
break
# 将opencv读取的图片进行预处理后，送网络中
# opencv 读取的图像s缩放为32X32的大小
resized_img = cv2.resize(image_raw, (32, 32), interpolation=cv2.INTER_LINEAR)
# bgr-->rgb
image_rgb = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)
# 转torch.Tensor
image = torch.from_numpy(image_rgb)
# 输入的图片是整数，要和网络参数(浮点数)保持一致
image = image.float()
# print(image.shape)
# [3,32,32] --> [1,3,32,32] 提升一张照片的维度
image = torch.reshape(image, [1, 3, 32, 32])
# 开始预测
model.eval()
with torch.no_grad():
output = model(image)
print(output)
idx = output.argmax(1)
reslut = class_id[idx]
print(reslut)
print(idx.item())
# 可视化
if idx >= 0:
cv2.putText(image_raw, reslut, (10, 20), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255))
print(reslut)
idx = -1
cv2.imshow("CIFAR", image_raw)
# 按下'q'键退出循环
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()
本篇博客演示了如何使用PyTorch构建一个简单的卷积神经网络来对CIFAR-10图像数据集进行分类。当然，这只是一个入门级别的示例，您可以通过调整网络结构、优化算法和超参数来进一步提高分类性能。通过不断实践和学习，您可以深入了解卷积神经网络以及如何在PyTorch中应用它们来解决实际的图像分类问题。
希望本篇博客对您理解基于PyTorch的图像分类卷积神经网络有所帮助！如果您有任何问题或建议，请随时在评论区留言。感谢阅读！
结束语
感谢你观看我的文章呐~本次航班到这里就结束啦
希望本篇文章有对你带来帮助 ，有学习到一点知识~
躲起来的星星也在努力发光，你也要努力加油（让我们一起努力叭）。
最后，博主要一下你们的三连呀（点赞、评论、收藏），不要钱的还是可以搞一搞的嘛~
不知道评论啥的，即使扣个666也是对博主的鼓舞吖 感谢