一、了解CIFAR-10数据集

CIFAR-10数据集5万张训练图像、1万张测试图像、10个类别、每个类别有6k个图像，图像大小32×32×3。下图列举了10个类，每一类随机展示了10张图片：
在这里插入图片描述
PyTorch 中的 torchvision.datasets 计算机视觉模块封装了 CIFAR10 数据集，如果需要使用可以直接导入。

导入代码：

from torchvision.datasets import CIFAR10

二、案例之导包

import torch
import torch.nn as nn
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor  # pip install torchvision -i https://mirrors.aliyun.com/pypi/simple/
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import matplotlib.pyplot as plt
from torchsummary import summary# 每批次样本数
BATCH_SIZE = 8

三、案例之创建数据集

# 1. 数据集基本信息
def create_dataset():# 加载数据集:训练集数据和测试数据# ToTensor: 将image（一个PIL.Image对象）转换为一个Tensortrain = CIFAR10(root='data', train=True, transform=ToTensor())valid = CIFAR10(root='data', train=False, transform=ToTensor())# 返回数据集结果return train, validif __name__ == '__main__':# 数据集加载train_dataset, valid_dataset = create_dataset()# 数据集类别print("数据集类别:", train_dataset.class_to_idx)# 数据集中的图像数据print("训练集数据集:", train_dataset.data.shape)print("测试集数据集:", valid_dataset.data.shape)# 图像展示plt.figure(figsize=(2, 2))plt.imshow(train_dataset.data[1])plt.title(train_dataset.targets[1])plt.show()

运行结果：

数据集类别: {'airplane': 0, 'automobile': 1, 'bird': 2, 'cat': 3, 'deer': 4, 'dog': 5, 'frog': 6, 'horse': 7, 'ship': 8, 'truck': 9}
训练集数据集: (50000, 32, 32, 3)
测试集数据集: (10000, 32, 32, 3)

图像：

在这里插入图片描述

四、案例之搭建神经网络（模型构建）

需要搭建的CNN网络结构如下:
在这里插入图片描述
我们要搭建的网络结构如下:

输入形状: 32x32；
第一个卷积层输入 3 个 Channel, 输出 6 个 Channel, Kernel Size 为: 3x3；
第一个池化层输入 30x30, 输出 15x15, Kernel Size 为: 2x2, Stride 为: 2；
第二个卷积层输入 6 个 Channel, 输出 16 个 Channel, Kernel Size 为 3x3；
第二个池化层输入 13x13, 输出 6x6, Kernel Size 为: 2x2, Stride 为: 2；
第一个全连接层输入 576 维, 输出 120 维；
第二个全连接层输入 120 维, 输出 84 维；
最后的输出层输入 84 维, 输出 10 维。

我们在每个卷积计算之后应用 relu 激活函数来给网络增加非线性因素。

# 模型构建
class ImageClassification(nn.Module):# 定义网络结构def __init__(self):super(ImageClassification, self).__init__()# 定义网络层：卷积层+池化层# 第一个卷积层, 输入图像为3通道,输出特征图为6通道,卷积核3*3self.conv1 = nn.Conv2d(3, 6, stride=1, kernel_size=3)# 第一个池化层, 核宽高2*2self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)# 第二个卷积层, 输入图像为6通道,输出特征图为16通道,卷积核3*3self.conv2 = nn.Conv2d(6, 16, stride=1, kernel_size=3)# 第二个池化层, 核宽高2*2self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)# 全连接层# 第一个隐藏层 输入特征576个(一张图像为16*6*6), 输出特征120个self.linear1 = nn.Linear(576, 120)# 第二个隐藏层self.linear2 = nn.Linear(120, 84)# 输出层self.out = nn.Linear(84, 10)# 定义前向传播def forward(self, x):# 卷积+relu+池化x = torch.relu(self.conv1(x))x = self.pool1(x)# 卷积+relu+池化x = torch.relu(self.conv2(x))x = self.pool2(x)# 将特征图做成以为向量的形式：相当于特征向量 全连接层只能接收二维数据集# 由于最后一个批次可能不够8，所以需要根据批次数量来改变形状# x[8, 16, 6, 6] --> [8, 576] -->8个样本,576个特征# x.size(0): 第1个值是样本数 行数# -1：第2个值由原始x剩余3个维度值相乘计算得到 列数(特征个数)x = x.reshape(x.size(0), -1)# 全连接层x = torch.relu(self.linear1(x))x = torch.relu(self.linear2(x))# 返回输出结果return self.out(x)if __name__ == '__main__':# 模型实例化model = ImageClassification()summary(model, input_size=(3,32,32), batch_size=1)

运行结果：

在这里插入图片描述

五、案例之编写训练函数（训练模型）

在训练时，使用多分类交叉熵损失函数，Adam 优化器。具体实现代码如下:

def train(model, train_dataset):# 构建数据加载器dataloader = DataLoader(train_dataset, batch_size=10, shuffle=True)criterion = nn.CrossEntropyLoss() # 构建损失函数optimizer = optim.Adam(model.parameters(), lr=1e-3) # 构建优化方法epoch = 100  # 训练轮数for epoch_idx in range(epoch):sum_num = 0   # 样本数量total_loss = 0.0  # 损失总和correct = 0  # 预测正确样本数start = time.time()  # 开始时间# 遍历数据进行网络训练for x, y in dataloader:model.train()output = model(x)loss = criterion(output, y)  # 计算损失optimizer.zero_grad()  # 梯度清零loss.backward()  # 反向传播optimizer.step()  # 参数更新correct += (torch.argmax(output, dim=-1) == y).sum()  # 计算预测正确样本数# 计算每次训练模型的总损失值 loss是每批样本平均损失值total_loss += loss.item()*len(y)  # 统计损失和sum_num += len(y)print('epoch:%2s loss:%.5f acc:%.2f time:%.2fs' %(epoch_idx + 1,total_loss / sum_num,correct / sum_num,time.time() - start))# 模型保存torch.save(model.state_dict(), 'model/image_classification.pth')#联合上面代码一起运行本代码
if __name__ == '__main__':# 数据集加载train_dataset, valid_dataset = create_dataset()# 模型实例化model = ImageClassification()# 模型训练train(model,train_dataset)

运行结果：

epoch: 1 loss:1.67102 acc:0.38 time:26.23s
epoch: 2 loss:1.35650 acc:0.51 time:27.63s
epoch: 3 loss:1.22355 acc:0.57 time:31.10s
epoch: 4 loss:1.14639 acc:0.59 time:66.37s
epoch: 5 loss:1.09468 acc:0.61 time:40.38s
。。。。。。

六、案例之编写预测函数（模型测试）

当已经训练好模型（model），并保存了模型参数（model.state_dict()），可直接实例化模型，并加载训练好的模型参数，然后对测试集中的1万条样本进行预测，查看模型在测试集上的准确率。

def eval(valid_dataset):# 构建数据加载器dataloader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False)# 加载模型并加载训练好的权重model = ImageClassification()model.load_state_dict(torch.load('model/image_classification.pth'))# 模型切换评估模式, 如果网络模型中有dropout/BN等层, 评估阶段不进行相应操作model.eval()# 计算精度total_correct = 0total_samples = 0# 遍历每个batch的数据，获取预测结果，计算精度for x, y in dataloader:output = model(x)total_correct += (torch.argmax(output, dim=-1) == y).sum()total_samples += len(y)# 打印精度print('Acc: %.2f' % (total_correct / total_samples))if __name__ == '__main__':train_dataset, valid_dataset = create_dataset()eval(valid_dataset)