PyTorch项目笔记（二）CIFAR-10图像分类

最新推荐文章于 2026-04-23 15:21:28 发布

原创

最新推荐文章于 2026-04-23 15:21:28 发布 · 5.8k 阅读

标签

#pytorch #分类 #python

本文介绍了基于CIFAR10数据集的图像分类任务，包括数据加载、使用Pytorch构建CNN模型、应用GPU加速训练过程以及在测试集上评估模型性能。通过随机旋转、翻转等图像增强方法处理训练数据，定义简单的卷积神经网络并利用CUDA进行GPU训练，最终达到59%的整体测试准确率。

本文将在CIFAR10的基础上介绍如何处理数据、图像分类以及使用GPU加速模型训练。

1 CIFAR-10数据集介绍

CIFAR-10数据集共有60000张尺寸为3 × 32 × 32的图像，包含以下10个类别：

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

其中训练集图像50000张，测试集图像10000张。

2 如何加载数据集

我们使用Pytorch提供的视觉工具包torchvision加载CIFAR-10数据集。在加载数据集时，我们可以对训练集和测试集分别进行预处理操作，此处需注意只有训练集需要进行图像增强。此外，使用torchvision包可以快速构建出DataLoader对象。具体代码如下：

#导入所需的包
import torch
import torchvision
import torchvision.transforms as transforms

#训练集图像预处理：
#对测试集进行图像增强可缓解过拟合的情况
train_transform = transforms.Compose([transforms.RandomRotation(degrees=10), #随机旋转
                                      transforms.RandomHorizontalFlip(p=0.5), #随机水平翻转
                                      transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1), #改变图像属性
                                      transforms.ToTensor(), #转Tensor
                                      transforms.Normalize( 
                                        mean=[0.5, 0.5, 0.5],
                                        std=[0.5, 0.5, 0.5]) #归一化
                                     ])

# 测试集图像预处理
test_transform = transforms.Compose([transforms.ToTensor(), #转Tensor
                                     transforms.Normalize(
                                        mean=[0.5, 0.5, 0.5],
                                        std=[0.5, 0.5, 0.5]) #归一化
                                    ])

#加载训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=train_transform)
#为训练集创建DataLoader对象
trainloader = torch.utils.data.DataLoader(trainset