基于 PyTorch 的 CIFAR-10 图像分类学习总结

在本次学习中,我通过 PyTorch 实现了一个基于 CNN 的 CIFAR-10 图像分类模型,完整掌握了从数据加载、模型构建到训练评估的全流程。以下是具体学习内容总结,包含关键代码实现:

CIFAR-10图像分类完整代码

V1

创建时间:15:53

关键步骤解析

1. 数据处理

数据处理是深度学习任务的基础,主要包括:

  • 数据转换:使用transforms将图像转为张量并归一化,使模型更容易学习
  • 数据集加载:利用torchvision.datasets加载 CIFAR-10 数据集
  • 数据加载器:通过DataLoader实现批处理、打乱数据和多进程加载
  • 数据可视化:编写imshow函数直观查看数据,验证数据加载是否正确

2. 模型构建

CNN 模型是处理图像任务的有效工具,本模型结构包括:

  • 卷积层:使用nn.Conv2d提取图像特征,通过卷积核捕获局部特征
  • 池化层:使用nn.MaxPool2d降低特征图维度,减少计算量并增强鲁棒性
  • 全连接层:使用nn.Linear实现最终分类,将提取的特征映射到 10 个类别
  • 激活函数:使用 ReLU 增加模型非线性表达能力,解决梯度消失问题

3. 模型训练

训练过程是模型学习的核心,主要步骤包括:

  • 损失函数:选择交叉熵损失函数,适合多分类任务
  • 优化器:使用 SGD 优化器,通过学习率和动量控制参数更新
  • 训练循环:多轮迭代训练,每个批次包括前向传播、损失计算、反向传播和参数更新
  • 设备加速:自动检测 GPU 并利用 CUDA 加速训练过程

4. 模型评估

通过测试集验证模型性能:

  • 加载测试数据并可视化
  • 使用训练好的模型进行预测
  • 对比预测结果与真实标签,直观评估模型分类效果

学习心得

  1. 数据预处理对模型性能影响很大,合适的归一化能加速模型收敛
  2. 网络结构设计需要平衡复杂度和计算效率,过深或过浅的网络都可能影响性能
  3. 训练过程中的超参数(如学习率、批次大小、训练轮次)需要根据实际情况调整
  4. GPU 加速能显著提高训练速度,特别是对于图像等数据量大的任务
  5. 可视化是调试和理解模型的有效手段,有助于发现数据或模型中的问题

通过这个实例,我掌握了 PyTorch 的基本使用方法和 CNN 图像分类的完整流程,为后续更复杂的深度学习任务打下了基础。

基于PyTorch的CIFAR-10图像分类实现

一、准备工作:导入必要的库

import torch

import torchvision

import torchvision.transforms as transforms

import matplotlib.pyplot as plt

import numpy as np

import torch.nn as nn

import torch.nn.functional as F

import torch.optim as optim

二、数据加载与预处理

1. 定义数据转换:将图像转为张量并归一化

transform = transforms.Compose(

[transforms.ToTensor(), # 转换为PyTorch张量

transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] # 归一化到-1, 1范围

)

2. 加载训练集和测试集

trainset = torchvision.datasets.CIFAR10(

root='./data', # 数据存储路径

train=True, # 训练集

download=True, # 如果本地没有数据则下载

transform=transform # 应用数据转换

)

testset = torchvision.datasets.CIFAR10(

root='./data',

train=False, # 测试集

download=True,

transform=transform

)

3. 创建数据加载器

trainloader = torch.utils.data.DataLoader(

trainset,

batch_size=4, # 批处理大小

shuffle=True, # 训练时打乱数据顺序

num_workers=2 # 多进程加载数据

)

testloader = torch.utils.data.DataLoader(

testset,

batch_size=4,

shuffle=False, # 测试时不打乱顺序

num_workers=2

)

4. 定义类别标签

classes = ('plane', 'car', 'bird', 'cat',

'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

三、数据可视化

定义图像显示函数

def imshow(img):

img = img / 2 + 0.5 # 反归一化

npimg = img.numpy() # 转换为numpy数组

调整通道顺序:从(C, H, W)转为(H, W, C)

plt.imshow(np.transpose(npimg, (1, 2, 0)))

plt.show()

获取一些随机的训练图像

dataiter = iter(trainloader)

images, labels = next(dataiter)

显示图像

imshow(torchvision.utils.make_grid(images))

打印标签

print(' '.join(f'{classeslabels\[j]:5s}' for j in range(4)))

四、构建CNN模型

class CNNNet(nn.Module):

def init(self):

super(CNNNet, self).init()

第一个卷积层:3输入通道,16输出通道,5x5卷积核

self.conv1 = nn.Conv2d(3, 16, 5)

第一个池化层:2x2池化核,步长为2

self.pool = nn.MaxPool2d(2, 2)

第二个卷积层:16输入通道,36输出通道,3x3卷积核

self.conv2 = nn.Conv2d(16, 36, 3)

第一个全连接层

self.fc1 = nn.Linear(36 * 6 * 6, 128)

第二个全连接层(输出层,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, 36 * 6 * 6)

第一个全连接层->ReLU

x = F.relu(self.fc1(x))

输出层

x = self.fc2(x)

return x

实例化模型并移动到可用设备(GPU/CPU)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net = CNNNet()

net.to(device)

五、定义损失函数和优化器

criterion = nn.CrossEntropyLoss() # 交叉熵损失函数,适合多分类任务

optimizer = optim.SGD( # 随机梯度下降优化器

net.parameters(),

lr=0.001, # 学习率

momentum=0.9 # 动量参数

)

六、训练模型

for epoch in range(10): # 训练10个epoch

running_loss = 0.0

for i, data in enumerate(trainloader, 0):

获取输入数据和标签,并移动到设备

inputs, labels = data0.to(device), data1.to(device)

清零梯度

optimizer.zero_grad()

前向传播、计算损失、反向传播、参数更新

outputs = net(inputs)

loss = criterion(outputs, labels)

loss.backward()

optimizer.step()

打印训练状态

running_loss += loss.item()

if i % 2000 == 1999: # 每2000个批次打印一次

print(f'{epoch + 1}, {i + 1:5d} loss: {running_loss / 2000:.3f}')

running_loss = 0.0

print('Finished Training')

七、模型测试

在测试集上进行预测

dataiter = iter(testloader)

images, labels = next(dataiter)

显示测试图像

imshow(torchvision.utils.make_grid(images))

print('GroundTruth: ', ' '.join(f'{classeslabels\[j]:5s}' for j in range(4)))

进行预测

images, labels = images.to(device), labels.to(device)

outputs = net(images)

_, predicted = torch.max(outputs, 1) # 获取预测概率最大的类别

print('Predicted: ', ' '.join(f'{classespredicted\[j]:5s}' for j in range(4)))

相关推荐
曲径海盗41 分钟前
【Python 学习第 32 天】字符串基础
python·学习笔记·编程入门
田里的水稻1 小时前
PyTorch_torchvision---数据集下载\TensorBoard可视化
人工智能·pytorch·python·深度学习·机器学习
程序员正茂1 小时前
Win11使用自定义训练集训练ncnn模型
python·yolo·图像识别·移动端·ncnn
zwd20051 小时前
Manim 画布尺寸与像素换算详解:frame_width、frame_height、pixel_width 和画质参数(0.21.0 实测)
python·分辨率·manim·数学动画·frame_width·manim教程
最强小杰2 小时前
Claude Opus 4.8 调用报 529 怎么办?并发限制触发过载——退避策略与降级方案完整指南
运维·网络·python·ai
databook2 小时前
MinMax vs Standard vs Robust:三种特征缩放怎么选
python·机器学习·数据分析
在世修行2 小时前
智学文枢:实体模型层
python·orm·sqlalchemy
张张123y2 小时前
# 后端第七课:FastAPI 实战入门(用刚学的 Python 写第一个 API 接口)
开发语言·python·mysql·api·fastapi
CCblandy2 小时前
一致性哈希到底省了多少搬迁?本机实测 3 轮:取模 6→7 搬 85.8%,V=128 只搬 14.3%,虚拟节点太少会被倾斜骗
后端·python·云原生·性能优化·架构
Yan-英杰2 小时前
【论文精读】 基于 TensorFlow 对卷积神经网络的优化研究:异构式 CPU+GPU 协同计算 + ReLU-Swish 激活函数
python·卷积神经网络