基于Python的人工智能应用案例系列（12）：CIFAR-10图像分类CNN

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，卷积核大小3x3
        self.conv2 = nn.Conv2d(6, 16, 3, 1) # 输入通道6，输出通道16
        self.fc1 = nn.Linear(6*6*16, 120)   # 全连接层1
        self.fc2 = nn.Linear(120, 84)       # 全连接层2
        self.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 = 0

with 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快速实现这些应用。

如果你觉得这篇博文对你有帮助，请点赞、收藏、关注我，并且可以打赏支持我！

欢迎关注我的后续博文，我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。

谢谢大家的支持！