Pytorch之CIFAR10分类卷积神经网络

前言
在计算机视觉领域,图像分类是一个重要的问题 。CIFAR-10是一个常用的图像分类数据集,包含10个类别的60000张32x32彩色图像 。在本篇博客中,我们将使用构建一个卷积神经网络(,CNN)来对CIFAR-10数据集中的图像进行分类 。
一、CIFAR-10数据集
CIFAR-10数据集包含10个类别,每个类别有6000张32x32大小的彩色图像 。这些类别分别是飞机、汽车、鸟类、猫、鹿、狗、青蛙、马、船和卡车 。数据集被分为训练集和测试集,每个集合各有50000张图像和10000张图像 。
二、实现卷积神经网络
我们将构建一个简单的卷积神经网络来进行图像分类 。以下是网络的架构:
卷积层( Layer):使用5x5的卷积核进行特征提取 。池化层( Layer):使用2x2的最大池化操作来减小特征图的大小 。全连接层(FullyLayer):将池化后的特征图展平并连接到一个全连接层,即上图中绿色的线条, 用于分类处理 。输出层:得到每个类别的概率分布 。
根据神经网络结构图,搭建网络模型时,需要计算的是卷积层中的和 。根据官方提供的计算公式,如下:
以第一层卷积层为例子,= 3,=32; Hin=32,Win=32;Hout=32,Wout=32; =5 。
注意,当//为int时,其sride[0]=[1]=int,[0]=[1]=int,[0]=[1]=int 。
将已知量带入公式,可计算得到=1,=2 。
以下是网络的代码实现:
from torch import nnimport 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 xif __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数据集并进行训练 。我们将使用的数据加载工具来处理数据集 。
# 准备数据集train_data = http://www.kingceram.com/post/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.01optimizer = 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 = dataoutputs = net(imgs)loss = loss_fn(outputs, targets)# 优化器优化模型optimizer.zero_grad()loss.backward()optimizer.step()total_train_step += 1if total_train_step % 100 == 0:print("训练次数:{},Loss:{}".format(total_train_step, loss.item()))writer.add_scalar("train_loss", loss.item(), total_train_step)