FA_CNN卷积神经网络十二_CIFAR10全流程PyTorch

FA ：Formulas and Algorithm，CNN：Convolutional Neural Network

1. 项目概述

本文基于 PyTorch 框架，使用经典的 CIFAR10 数据集，完整演示一个图像分类模型的训练与验证流程。内容涵盖：

• 使用 GPU 加速训练一个自定义卷积神经网络（CNN）
• 训练 30 个轮次（epoch），并记录训练过程中的损失与准确率
• 训练完成后，使用训练好的模型对真实图片（狗、飞机）进行识别验证

通过本文，你将掌握 PyTorch 从数据加载、模型搭建、训练、测试到模型保存与加载的完整套路。

2. 环境准备

在开始之前，请确保你的环境中已安装以下依赖：

pip install torch torchvision pillow tensorboard

• torch / torchvision：PyTorch 核心库与视觉工具库
• pillow：图像处理库，用于读取和预处理图片
• tensorboard：训练过程可视化工具

同时，建议准备一块可用的 NVIDIA GPU（显存建议 4GB 以上），以加速训练过程。

3. 数据集准备

CIFAR10 数据集包含 10 个类别的 60000 张 32×32 彩色图片，其中训练集 50000 张，测试集 10000 张。类别包括飞机、汽车、鸟、猫、鹿、狗、青蛙、马、船、卡车。

使用 torchvision.datasets.CIFAR10 可以自动下载并加载数据集：

import torchvision

# 训练集
train_data = torchvision.datasets.CIFAR10(
    "./dataset",
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=True
)

# 测试集
test_data = torchvision.datasets.CIFAR10(
    "./dataset",
    train=False,
    transform=torchvision.transforms.ToTensor(),
    download=True
)

print("训练数据集的长度：{}".format(len(train_data)))
print("测试数据集的长度：{}".format(len(test_data)))

输出结果：

训练数据集的长度：50000
测试数据集的长度：10000

4. 构建卷积神经网络模型

我们定义一个名为 Tudui 的卷积神经网络，结构如下：

• 3 层卷积层（Conv2d），每层后接最大池化（MaxPool2d）
• 展平层（Flatten）
• 2 层全连接层（Linear）

from torch import nn

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),   # 输入通道3，输出通道32，卷积核5×5，步长1，填充2
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),                # 展平后变成 64*4*4
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.model1(x)
        return x

5. 训练配置

5.1 设备选择

优先使用 GPU 进行训练，并提供了多种设备指定方式：

import torch

# 方式一：直接指定 cuda
device = torch.device("cuda")

# 方式二：指定具体 GPU 编号
# device = torch.device("cuda:0")

# 方式三：自动判断
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

5.2 数据加载器

使用 DataLoader 将数据集分批加载，每批 64 张图片：

from torch.utils.data import DataLoader

train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

5.3 损失函数与优化器

• 损失函数：交叉熵损失（CrossEntropyLoss），适用于多分类任务
• 优化器：随机梯度下降（SGD），学习率设为 0.01

# 创建网络模型并移动到设备
tudui = Tudui().to(device)

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

# 优化器
learning = 0.01  # 1e-2 就是 0.01
optimizer = torch.optim.SGD(tudui.parameters(), learning)

6. 完整训练流程

6.1 训练主循环

训练共进行 30 个轮次，每轮包含训练步骤和测试步骤：

import time
from torch.utils.tensorboard import SummaryWriter

# 记录训练/测试次数
total_train_step = 0
total_test_step = 0
epoch = 30

# TensorBoard 记录
writer = SummaryWriter("logs")
start_time = time.time()

for i in range(epoch):
    print("-----第 {} 轮训练开始-----".format(i + 1))

    # 训练步骤
    tudui.train()  # 启用 dropout、batchnorm 等层
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)

        outputs = tudui(imgs)
        loss = loss_fn(outputs, targets)

        # 优化器调优
        optimizer.zero_grad()   # 梯度清零
        loss.backward()         # 反向传播
        optimizer.step()        # 更新参数

        total_train_step += 1
        if total_train_step % 100 == 0:
            end_time = time.time()
            print("训练次数：{}，Loss：{}".format(total_train_step, loss.item()))
            print("耗时：{} 秒".format(end_time - start_time))

    writer.add_scalar("train_loss", loss.item(), total_train_step)

    # 测试步骤
    tudui.eval()  # 关闭 dropout、batchnorm 等层
    total_test_loss = 0
    total_accuracy = 0

    with torch.no_grad():  # 不计算梯度
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)

            outputs = tudui(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss += loss.item()

            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_accuracy", total_accuracy / test_data_size, total_test_step)
    total_test_step += 1

    # 保存模型
    torch.save(tudui, "./model/tudui_{}.pth".format(i))
    print("模型已保存")

writer.close()

6.2 训练过程关键输出

以下是训练过程中的部分关键输出，展示了损失下降与准确率提升的趋势：

训练数据集的长度：50000
测试数据集的长度：10000
-----第 1 轮训练开始-----
训练次数：100，Loss：2.2807211875915527
整体测试集上的Loss：305.39513182640076
整体测试集上的正确率：0.29249998927116394
模型已保存
-----第 10 轮训练开始-----
训练次数：7100，Loss：1.284026861190796
整体测试集上的Loss：195.0189254283905
整体测试集上的正确率：0.560699999332428
模型已保存
-----第 20 轮训练开始-----
训练次数：14900，Loss：0.6770450472831726
整体测试集上的Loss：168.21677833795547
整体测试集上的正确率：0.6384999752044678
模型已保存
-----第 30 轮训练开始-----
训练次数：22700，Loss：0.733674168586731
整体测试集上的Loss：174.22026354074478
整体测试集上的正确率：0.6488999724388123
模型已保存

从输出可以看出：

• 训练损失从第 1 轮的约 2.28 下降到第 30 轮的约 0.73
• 测试集准确率从约 29.2% 提升到约 64.9%
• 训练耗时约 317 秒（约 5.3 分钟）

7. 模型验证：识别真实图片

训练完成后，我们使用保存的模型对真实图片进行识别。这里以「狗」和「飞机」两张图片为例。

7.1 验证狗图片

import torchvision
from PIL import Image
from torch import nn
import torch

image_path = "imgs/dog.png"
image = Image.open(image_path)          # PIL 类型图片
image = image.convert("RGB")            # RGBA 转 RGB
print(image)

# 图片预处理：缩放至 32×32 并转为张量
transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),
    torchvision.transforms.ToTensor()
])
image = transform(image)
print(image.shape)

# 加载模型（GPU 训练，CPU 加载）
model = torch.load("model/tudui_29.pth", map_location=torch.device('cpu'))
print(model)

# 转为四维，符合网络输入要求
image = torch.reshape(image, (1, 3, 32, 32))

model.eval()
with torch.no_grad():
    output = model(image)

print(output)
print(output.argmax(1))  # 概率最大类别的输出

输出结果：

<PIL.Image.Image image mode=RGB size=307x173 at 0x1A23417EDA0>
torch.Size([3, 32, 32])
tensor([[ 1.3921, -5.8952,  6.7765,  2.6106,  3.3179,  6.3351, -5.1229,  3.0367,
         -8.6725, -5.5459]], grad_fn=<AddmmBackward0>)
tensor([2])

输出 tensor([2]) 表示模型预测该图片属于类别 2，即「鸟」。注意：这里模型将狗识别为鸟，说明模型在部分类别上仍存在误判，可通过增加训练轮次或优化模型结构来提升准确率。

7.2 验证飞机图片

import torchvision
from PIL import Image
from torch import nn
import torch

image_path = "imgs/plane.png"
image = Image.open(image_path)
image = image.convert("RGB")
print(image)

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),
    torchvision.transforms.ToTensor()
])
image = transform(image)
print(image.shape)

# 加载模型
model = torch.load("model/tudui_29.pth", map_location=torch.device('cpu'))
print(model)

image = torch.reshape(image, (1, 3, 32, 32))

model.eval()
with torch.no_grad():
    output = model(image)

print(output)
print(output.argmax(1))

输出结果：

<PIL.Image.Image image mode=RGB size=245x181 at 0x1A236055908>
torch.Size([3, 32, 32])
tensor([[ 3.9218,  1.2536,  1.1701,  0.1064,  2.2779, -4.4623,  0.0393, -4.2867,
          2.3805, -2.1321]], grad_fn=<AddmmBackward0>)
tensor([0])

输出 tensor([0]) 表示模型预测该图片属于类别 0，即「飞机」。识别正确 ！

8. 总结与优化建议

8.1 总结

本文完整演示了使用 PyTorch 训练 CIFAR10 图像分类模型的流程：

1. 数据准备 ：使用 torchvision.datasets.CIFAR10 自动下载并加载数据
2. 模型构建：自定义 CNN 网络，包含卷积、池化、全连接层
3. 训练配置：交叉熵损失 + SGD 优化器，GPU 加速
4. 训练与测试：30 轮训练，每轮记录损失与准确率
5. 模型保存与加载 ：使用 torch.save / torch.load 保存和加载模型
6. 图片验证：对真实图片进行预处理并预测类别

8.2 优化建议

• 增加训练轮次：当前 30 轮准确率约 64.9%，可尝试训练 50 轮或更多
• 调整学习率 ：可尝试使用学习率衰减策略，如 StepLR
• 数据增强：使用随机裁剪、翻转等操作提升模型泛化能力
• 更换优化器 ：尝试 Adam 优化器，通常收敛更快
• 模型结构优化：增加卷积层深度或使用残差结构（ResNet）

9. 完整代码

以下是训练部分的完整代码，可直接复制运行：

import torchvision
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time

# 设备选择
device = torch.device("cuda")

# 定义模型
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.model1(x)
        return x

# 准备数据集
train_data = torchvision.datasets.CIFAR10(
    "./dataset", train=True,
    transform=torchvision.transforms.ToTensor(), download=True
)
test_data = torchvision.datasets.CIFAR10(
    "./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))

# 数据加载器
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

# 创建模型
tudui = Tudui().to(device)

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

# 优化器
learning = 0.01
optimizer = torch.optim.SGD(tudui.parameters(), learning)

# 训练参数
total_train_step = 0
total_test_step = 0
epoch = 30

# TensorBoard
writer = SummaryWriter("logs")
start_time = time.time()

for i in range(epoch):
    print("-----第 {} 轮训练开始-----".format(i + 1))

    tudui.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)

        outputs = tudui(imgs)
        loss = loss_fn(outputs, targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        if total_train_step % 100 == 0:
            end_time = time.time()
            print("训练次数：{}，Loss：{}".format(total_train_step, loss.item()))
            print("耗时：{} 秒".format(end_time - start_time))

    writer.add_scalar("train_loss", loss.item(), total_train_step)

    tudui.eval()
    total_test_loss = 0
    total_accuracy = 0

    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)

            outputs = tudui(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss += loss.item()

            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_accuracy", total_accuracy / test_data_size, total_test_step)
    total_test_step += 1

    torch.save(tudui, "./model/tudui_{}.pth".format(i))
    print("模型已保存")

writer.close()