CIFAR-10是一个非常经典的图像分类数据集，类似于MNIST，但每个图像有三个颜色通道（RGB），而不是一个灰度通道。每张CIFAR-10图像的大小为 (3,32,32)，即32x32像素，具有3个颜色通道。该数据集包含10类图像，每类各5000张训练图像，总共50000张训练图像和10000张测试图像。本文将使用卷积神经网络（CNN）对该数据集进行分类。
1. 数据加载
我们使用PyTorch的 torchvision 模块来加载CIFAR-10数据集，并将其转换为张量格式。第一次运行时，数据会自动下载到指定路径。
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader# 定义转换
transform = transforms.ToTensor()# 加载训练集和测试集
train_data = datasets.CIFAR10(root='../data', train=True, download=True, transform=transform)
test_data = datasets.CIFAR10(root='../data', train=False, download=True, transform=transform)# 定义数据加载器
train_loader = DataLoader(train_data, batch_size=10, shuffle=True)
test_loader = DataLoader(test_data, batch_size=10, shuffle=False)
2. 数据探索与可视化
为了直观地了解数据集的内容，我们将显示一个批次的图像并输出其对应的标签。
import matplotlib.pyplot as plt
from torchvision.utils import make_grid# 获取一个批次的图像和标签
for images, labels in train_loader:break# 显示图像
im = make_grid(images, nrow=5)
plt.figure(figsize=(10,4))
plt.imshow(np.transpose(im.numpy(), (1, 2, 0)))
3. 模型构建
我们使用卷积神经网络（CNN）来构建模型。该网络包括两个卷积层和三个全连接层。
import torch.nn as nn
import torch.nn.functional as F# 定义卷积神经网络模型
class ConvolutionalNetwork(nn.Module):def __init__(self):super().__init__()self.conv1 = nn.Conv2d(3, 6, 3, 1) # 输入通道3，输出通道6，卷积核大小3x3self.conv2 = nn.Conv2d(6, 16, 3, 1) # 输入通道6，输出通道16self.fc1 = nn.Linear(6*6*16, 120) # 全连接层1self.fc2 = nn.Linear(120, 84) # 全连接层2self.fc3 = nn.Linear(84, 10) # 全连接层3，输出10类def forward(self, X):X = F.relu(self.conv1(X)) # 第一卷积层 + 激活函数X = F.max_pool2d(X, 2, 2) # 池化层X = F.relu(self.conv2(X)) # 第二卷积层 + 激活函数X = F.max_pool2d(X, 2, 2) # 池化层X = X.view(-1, 6*6*16) # 展平X = F.relu(self.fc1(X)) # 全连接层1 + 激活函数X = F.relu(self.fc2(X)) # 全连接层2 + 激活函数X = self.fc3(X) # 输出层return X
4. 模型训练
我们定义损失函数和优化器，并对模型进行训练。
import torch.optim as optim# 实例化模型、损失函数和优化器
model = ConvolutionalNetwork()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练模型
epochs = 10
for i in range(epochs):for X_train, y_train in train_loader:optimizer.zero_grad()y_pred = model(X_train)loss = criterion(y_pred, y_train)loss.backward()optimizer.step()
5. 模型评估
在训练完成后，我们可以通过在测试集上进行预测来评估模型的性能。
correct = 0
total = 0with torch.no_grad():for X_test, y_test in test_loader:y_val = model(X_test)predicted = torch.max(y_val, 1)[1]total += y_test.size(0)correct += (predicted == y_test).sum().item()print(f'Test accuracy: {correct / total * 100:.2f}%')
6. 混淆矩阵
为了更深入地分析模型的性能，我们可以生成混淆矩阵来观察哪些类别被误分类。
from sklearn.metrics import confusion_matrix
import seaborn as sn
import pandas as pd# 生成混淆矩阵
y_true = []
y_pred = []with torch.no_grad():for X_test, y_test in test_loader:outputs = model(X_test)_, predicted = torch.max(outputs, 1)y_true.extend(y_test.cpu().numpy())y_pred.extend(predicted.cpu().numpy())cm = confusion_matrix(y_true, y_pred)
df_cm = pd.DataFrame(cm, index=[i for i in range(10)], columns=[i for i in range(10)])
plt.figure(figsize=(10,7))
sn.heatmap(df_cm, annot=True)
结语
在本文中，我们通过卷积神经网络（CNN）成功地对CIFAR-10数据集中的图像进行了分类。我们从数据加载和预处理开始，构建了一个包含卷积层和全连接层的深度学习模型。通过对模型的训练和测试，我们不仅评估了模型的准确性，还使用混淆矩阵深入分析了分类效果。
CIFAR-10数据集比MNIST复杂得多，因为它包含彩色图像且图像尺寸较大，类别之间的差异也更加多样化。这使得分类任务更具挑战性，但通过设计合适的网络结构，我们能够达到不错的效果。
在未来的应用中，进一步优化模型的参数、增加数据增强、或者采用迁移学习等技术，可以帮助我们提升模型的表现。在更复杂的图像分类任务中，这些方法将变得更加关键。
通过这个案例，我们看到了深度学习在图像分类中的强大能力，以及如何使用PyTorch快速实现这些应用。
如果你觉得这篇博文对你有帮助，请点赞、收藏、关注我，并且可以打赏支持我！
欢迎关注我的后续博文，我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。
谢谢大家的支持！
