基于Pytorch框架的CIFAR-10图像分类任务（附带完整代码）

        本文主要实现在pytorch框架下，训练CIFAR数据集，通过观察训练和验证的误差、准确率图像来进一步改善。保存最好的模型。测试集打印整体准确率和每一类别的准确率，并生成混淆矩阵，将其中每一个错误的图片并保存下来。

语言：python

实现方式：pytorch框架, CPU

关键词: CIFAR-10数据集、Dataset和Dataloader、SummaryWriter画图、网络模型搭建、混淆矩阵、统计所有错误类别、Adam优化器

前言：CIFAR-10数据集介绍

        CIFAR-10是一个包含60000张32x32像素的RGB图像的数据集，共有10个类别。其中有50000张用于训练集，10000张用于测试集。每个图像都属于10个类别之一：飞机、汽车、鸟、猫、鹿、狗、青蛙、马、船和卡车。其中训练集和测试集是互相独立的。

        CIFAR-10数据集是机器学习中一个非常常用的图像分类数据集之一，也是深度学习算法的基准测试数据集之一。当我们想要测试新的算法或新的网络结构时，可以使用CIFAR-10数据集验证它们的性能。由于数据集的图像分辨率比较低，因此我们可以较为快速地进行模型训练，并能够在低成本的计算资源上进行模型开发和优化。

        当我们将CIFAR-10数据集导入到计算机中，我们可以使用各种工具和框架来预处理、建模和训练我们的分类器，例如使用Python中的TensorFlow和PyTorch等深度学习框架。（如何下载数据集请看下面代码）

        如下图所示：

一、搭建网络模型

        1.代码展示:

import torch
from torch import nn
from torch.nn import ReLU,Linear,Flatten,Sequential
from torch.nn import MaxPool2d
from torch.nn import Conv2d

#搭建网络模型
class Yangc(nn.Module):
    def __init__(self):
        super(Yangc,self).__init__()
        self.model=Sequential(
            Conv2d(3, 32, 5, padding = 2),#有公式可以计算
            ReLU(),
            MaxPool2d(2),    
   
            Conv2d(32, 32, 5, padding = 2),
            ReLU(),
            MaxPool2d(2),

            Conv2d(32, 64, 5, padding = 2),
            ReLU(),
            MaxPool2d(2),

            Flatten(),#展平
            Linear(1024, 512),
            Linear(512, 64),
            Linear(64, 10)
        )

    def forward(self,x):
        x = self.model(x)
        return x
    
if __name__ == '__main__':#写主函数，是为了检验神经网络模型是否搭建正确
    yangc = Yangc()
    input = torch.ones(64,3,32,32)
    output = yangc(input)
    print(output.shape)

输出:

torch.Size([64, 10])

二、开始训练、验证

1.准备数据集

import torch
import os
from torch import nn
from torch.utils.data import DataLoader
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from sklearn.model_selection import train_test_split
import time#用于计时
from model import Yangc#导入模型

#准备数据集，要归一化处理
train_data = torchvision.datasets.CIFAR10(root = "/home/yc/vs_code/model_neural_network/dataset_CIFAR-10",
                                        train = True,#表示为训练集
                                        transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
                                                                                transforms.Normalize(mean = [0.5, 0.5, 0.5],std = [0.5, 0.5, 0.5])
                                                                                ]),#将图片转为tensor格式，并归一化
                                        download = True#若没有CCIFAR-10数据集，则会自动下载并保存。若有数据集，则不会下载
                                        )


#划分数据集
train_data, val_data = train_test_split(train_data, train_size = 0.8, test_size = 0.2)#将训练集按8:2的比例划分为40000训练集，10000验证集

#打印数据集长度
print("训练数据集的长度为：{}".format(len(train_data)))
print("验证数据集的长度为：{}".format(len(val_data)))

# 利用 DataLoader 来加载数据集
train_dataloader = DataLoader(train_data, batch_size = 64, num_workers = 5)
val_dataloader = DataLoader(train_data, batch_size = 64, num_workers = 5)

#定义训练设备
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

说明：num_workers=5是基于当时的电脑设备而设置的，故不同的电脑配置，该值不一样，具体可参考这篇博客：num_workers的设置

输出:

训练数据集的长度为：40000
验证数据集的长度为：10000
Using cpu device

2.定义训练、验证函数

#画图
writer = SummaryWriter("img_loss_acc")

#搭建网络模型
model = Yangc()

# 损失函数
loss_fn = nn.CrossEntropyLoss()

# 优化器
learning_rate = 0.001
optimizer = torch.optim.Adam(model.parameters(),lr = learning_rate)

#定义训练函数
def train(train_loader, model, loss_fn, optimizer):
    corrects,num = 0.,0.#不加最后一个0就默认为浮点型
    total_loss = 0.
    model.train()

    for i,(X,targets) in enumerate(train_loader):
        outputs = model(X)
        optimizer.zero_grad()
        loss = loss_fn(outputs, targets)
        total_loss += loss.item()

        # 优化器优化模