PyTorch：学习 CIFAR-10 分类

yukoyu2025-04-18 15:22

🔍 开始你的图像分类之旅：一步一步学习 CIFAR-10 分类

图像分类是计算机视觉中最基础的任务之一，如果你是初学者，那么以 CIFAR-10 为训练场是一个不错的选择。本文一步一步带你从零开始，学习如何用深度学习模型实现图像分类。

一、CIFAR-10 数据集是什么？

CIFAR-10 是一个小型图像分类数据集，共包括 10 个类别：✈ 飞机（airplane）🚗 汽车（automobile）🐦 鸟（bird）🐱 猫（cat）🦌 鹿（deer）🐶 狗（dog）🐸 青蛙（frog）🐴 马（horse）🚢 船（ship）🚚 卡车（truck）

每张图片都是 32x32 的小图，有 RGB 三个颜色通道。

总共有 60000 张图，其中:

  • 训练集: 50000 张
  • 测试集: 10000 张

这些图片内容丰富，分辨率低，适合初学者练手。

二、模型训练的整体流程

我们用以下流程完成图像分类：

  1. 数据加载和预处理
  2. 构建模型（CNN）
  3. 设置损失函数和优化器
  4. 训练模型（前向 + 反向传播 + 更新参数）
  5. 测试模型效果

🧠 类比理解： 把整个过程比作"学会识别水果"：

  • 数据加载：收集不同水果的照片
  • 模型：像是大脑处理这些图像的神经元网络
  • 损失函数：告诉我们判断错误的严重程度
  • 优化器：帮助我们不断修正错误，直到准确

三、数据加载和预处理

我们使用 PyTorch 中的 transforms 来将图片：

  • 转换成 Tensor（张量）
  • 正则化颜色值到 -1~1 之间，加快模型收敛
python 复制代码 
import torch
import torchvision
import torchvision.transforms as transforms

# 定义图像转换操作：将图片转换为 Tensor，并进行标准化
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图片转为 Tensor 类型，方便计算
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 将图片标准化，均值和标准差都设为0.5
])

# 加载 CIFAR-10 训练数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
# 使用 DataLoader 进行批量加载数据，batch_size 是每次加载的图片数量
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
                                          shuffle=True)  # shuffle=True 表示打乱数据

# 加载 CIFAR-10 测试数据集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64,
                                         shuffle=False)  # 测试集不需要打乱数据

🎓 例子解释：

  • 如果 batch_size=64，就意味着每次训练取 64 张图片
  • shuffle=True 可以打乱图片顺序，防止模型记住顺序而不是学习特征

🍎 通俗类比：

  • ToTensor 就像是把一张照片转成表格（数字表示颜色）
  • Normalize 就像是统一标准，把所有颜色亮度调成统一区间，好比较

四、构建 CNN 模型

CNN（卷积神经网络）特别适合处理图像。我们构建一个简单 CNN 模型，包含：

  • 两个卷积层 + ReLU 激活
  • 两个最大池化层（缩小图片尺寸）
  • 一个全连接隐藏层 + 一个输出层（10类）
python 复制代码 
import torch.nn as nn
import torch.nn.functional as F

# 定义简单的卷积神经网络
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        
        # 第一层卷积层，输入通道为3（RGB图像），输出通道为32，卷积核大小为3x3，padding=1 保证输出尺寸不变
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        # 最大池化层：2x2池化，用来减少图像尺寸
        self.pool = nn.MaxPool2d(2, 2)
        
        # 第二层卷积层，输入通道为32，输出通道为64
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        
        # 全连接层：将卷积层输出展平为一维向量，连接到一个128维的隐藏层
        self.fc1 = nn.Linear(64 * 8 * 8, 128)  # CIFAR-10图像尺寸32x32，经过两次池化后尺寸为8x8
        # 输出层：10类，CIFAR-10数据集包含10个类别
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 第一层卷积 + ReLU 激活 + 池化
        x = self.pool(F.relu(self.conv1(x)))
        # 第二层卷积 + ReLU 激活 + 池化
        x = self.pool(F.relu(self.conv2(x)))
        
        # 展平特征图为一维向量，便于输入全连接层
        x = x.view(-1, 64 * 8 * 8)
        
        # 全连接层 + ReLU 激活
        x = F.relu(self.fc1(x))
        # 输出层，返回每个类别的预测概率
        x = self.fc2(x)
        return x

📷 类比：

  • 卷积操作就像"扫描照片"的滤镜，用来提取边缘、颜色块等图像特征
  • 最大池化像是"缩略图"，保留最显著的特征，减少计算量

五、设置损失函数和优化器

python 复制代码 
import torch.optim as optim

# 初始化模型并将其放到 GPU 或 CPU 上
model = SimpleCNN()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 损失函数：交叉熵损失函数，用于多分类问题
criterion = nn.CrossEntropyLoss()

# 优化器：Adam 优化器，适用于大部分情况，学习率设置为0.001
optimizer = optim.Adam(model.parameters(), lr=0.001)

🧠 例子类比：

  • 损失函数就像考试成绩：越高说明你错越多
  • 优化器就像"老师指出你的错误"并教你怎么改正

📘 数值示例：

  • 如果模型预测飞机的概率是 [0.1, 0.05, ..., 0.7]（第10类）
  • 但真实标签是第1类（飞机），交叉熵损失会很大
  • 优化器就会调整参数，使下一次飞机的概率尽量靠近第1类

六、训练模型：前向、反向、更新

python 复制代码 
# 训练过程：迭代多个 epoch，每个 epoch 会遍历所有训练数据
for epoch in range(10):  # 总共训练 10 个 epoch
    running_loss = 0.0  # 每个 epoch 初始化损失值为 0
    for inputs, labels in trainloader:  # 遍历每个 batch
        inputs, labels = inputs.to(device), labels.to(device)  # 将数据移到 GPU 或 CPU

        optimizer.zero_grad()  # 清除上一次的梯度

        outputs = model(inputs)  # 前向传播，得到每张图片的预测结果
        loss = criterion(outputs, labels)  # 计算损失值

        loss.backward()  # 反向传播，计算梯度

        optimizer.step()  # 更新模型参数

        running_loss += loss.item()  # 累加损失值
    print(f"Epoch {epoch + 1}, Loss: {running_loss:.3f}")  # 打印当前 epoch 的损失

📐 具体数值例子：

  • 模型初始权重是 0.3，预测错误 → loss = 2.5
  • 反向传播算出权重梯度是 -0.8
  • 学习率为 0.01，更新后权重 = 0.3 - 0.01 × (-0.8) = 0.308

七、测试模型效果

python 复制代码 
correct = 0  # 初始化正确预测的个数
total = 0  # 初始化总预测的个数
model.eval()  # 设置模型为评估模式，关闭 Dropout 等训练时的特殊操作

with torch.no_grad():  # 在测试时，不需要计算梯度，减少计算量
    for data in testloader:  # 遍历测试集中的数据
        images, labels = data
        images, labels = images.to(device), labels.to(device)  # 将数据移到 GPU 或 CPU
        outputs = model(images)  # 获取模型的输出
        _, predicted = torch.max(outputs, 1)  # 获取预测结果，torch.max 返回最大值和其索引，这里我们只取索引

        total += labels.size(0)  # 累加总的测试样本数量
        correct += (predicted == labels).sum().item()  # 统计预测正确的样本数量

print(f"测试准确率: {100 * correct / total:.2f}%")  # 打印测试集上的准确率

📊 例子：

  • 如果 total=10000，correct=7000，准确率就是 70%

🏁 完整代码快速运行包

python 复制代码 
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import matplotlib.pyplot as plt  # 添加matplotlib用于可视化
from matplotlib import rcParams  # 用于设置字体

# 1. 数据加载与预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化到[-1, 1]之间
])

# 加载训练集与测试集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

# 2. 定义模型：简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        # 第一个卷积层：3个输入通道，32个输出通道，卷积核大小3x3，padding为1
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        # 最大池化层：2x2池化
        self.pool = nn.MaxPool2d(2, 2)
        # 第二个卷积层：32个输入通道，64个输出通道，卷积核大小3x3，padding为1
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        # 全连接层，输入大小64x8x8，输出128
        self.fc1 = nn.Linear(64 * 8 * 8, 128)
        # 最后一层全连接层，输出10个类别
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # 卷积+ReLU+池化
        x = self.pool(F.relu(self.conv2(x)))  # 卷积+ReLU+池化
        x = x.view(-1, 64 * 8 * 8)  # 展平数据，准备全连接
        x = F.relu(self.fc1(x))  # 全连接+ReLU
        x = self.fc2(x)  # 最后一层输出
        return x

# 3. 初始化模型、损失函数与优化器
model = SimpleCNN()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 判断是否使用GPU
model.to(device)  # 将模型转移到GPU或CPU上

criterion = nn.CrossEntropyLoss()  # 使用交叉熵作为损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器，学习率为0.001

# 4. 训练模型
for epoch in range(10):  # 训练10个epoch
    running_loss = 0.0
    for inputs, labels in trainloader:
        inputs, labels = inputs.to(device), labels.to(device)  # 数据转移到GPU或CPU

        optimizer.zero_grad()  # 清除上一次的梯度
        outputs = model(inputs)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 优化器更新参数

        running_loss += loss.item()  # 累加损失
    print(f"Epoch {epoch + 1}, Loss: {running_loss:.3f}")  # 打印每个epoch的损失

# 5. 测试模型效果
correct = 0
total = 0
model.eval()  # 将模型设置为评估模式

with torch.no_grad():  # 禁止计算梯度，提高效率
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)  # 数据转移到GPU或CPU
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)  # 获取最大概率的类别
        total += labels.size(0)  # 累加总样本数
        correct += (predicted == labels).sum().item()  # 统计正确的样本数

        # 可视化一个批次的预测结果
        plt.figure(figsize=(10, 10))
        for i in range(8):  # 显示前8张图片
            plt.subplot(2, 4, i + 1)
            plt.imshow(images[i].cpu().permute(1, 2, 0) * 0.5 + 0.5)  # 反归一化并显示图片
            plt.title(f"Label: {labels[i].item()}\nPrediction: {predicted[i].item()}")
            plt.axis('off')
        plt.show()
        break  # 仅显示一个批次

# 输出测试准确率
print(f"测试准确率: {100 * correct / total:.2f}%")

你可以将以下内容保存为 train_cifar10.py 并运行：

复制代码 
python train_cifar10.py

💡 不需要修改任何内容就能开始训练和测试！有 CUDA 就用 GPU，否则自动使用 CPU。

相关推荐
西岸行者5 天前
学习笔记：SKILLS 能帮助更好的vibe coding
笔记·学习
悠哉悠哉愿意5 天前
【单片机学习笔记】串口、超声波、NE555的同时使用
笔记·单片机·学习
别催小唐敲代码5 天前
嵌入式学习路线
学习
毛小茛5 天前
计算机系统概论——校验码
学习
babe小鑫5 天前
大专经济信息管理专业学习数据分析的必要性
学习·数据挖掘·数据分析
winfreedoms5 天前
ROS2知识大白话
笔记·学习·ros2
在这habit之下5 天前
Linux Virtual Server（LVS）学习总结
linux·学习·lvs
我想我不够好。5 天前
2026.2.25监控学习
学习
im_AMBER5 天前
Leetcode 127 删除有序数组中的重复项 | 删除有序数组中的重复项 II
数据结构·学习·算法·leetcode
CodeJourney_J5 天前
从“Hello World“ 开始 C++
c语言·c++·学习
热门推荐
01GitHub 镜像站点02OpenClaw 使用和管理 MCP 完全指南03OpenClaw + 飞书（Feishu）环境搭建指南04【OpenClaw 本地实战 Ep.3】突破瓶颈：强制修改 openclaw.json 解锁 32k 上下文记忆05Claude Code + GLM4.7 避坑指南：解决 Unable to connect to Anthropic services06OpenClaw优化飞书API 额度已耗尽问题07Window 10部署openclaw报错node.exe : npm error code 12808OpenClaw大龙虾机器人完整安装教程09Clawdbot部署教程：解决‘gateway token missing’授权问题的完整步骤10小黑课堂计算机二级WPSoffice题库软件下载安装教程（2026年3月最新版）