卷积神经网络（CNN）图像分类案例

一个基于卷积神经网络（Convolutional Neural Network, CNN）的图像分类完整案例。该案例使用 PyTorch 构建一个 CNN 模型来对 CIFAR-10 数据集进行分类。

1. 环境准备

安装 PyTorch 和必要的库（如未安装）：

pip install torch torchvision matplotlib

2. 导入必要的库

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

3. 数据加载与预处理

3.1 CIFAR-10 数据集

CIFAR-10 是一个常用的图像分类数据集，包含 10 个类别，每个类别有 6000 张 32x32 的彩色图片。

3.2 数据预处理

• 图像标准化：将像素值从 [0, 255] 映射到 [0, 1]，并进一步标准化为零均值和单位方差。

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化
])

# 加载训练集和测试集
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)

# CIFAR-10 类别
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

3.3 显示部分图片

# 显示部分训练图片
def imshow(img):
    img = img / 2 + 0.5  # 去标准化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

dataiter = iter(trainloader)
images, labels = next(dataiter)
imshow(torchvision.utils.make_grid(images[:4]))  # 显示前4张图片
print(' '.join(f'{classes[labels[j]]}' for j in range(4)))

4. 构建 CNN 模型

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 卷积层 1：输入通道 3，输出通道 32，卷积核大小 3x3
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        # 卷积层 2：输入通道 32，输出通道 64，卷积核大小 3x3
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        # 池化层：最大池化，窗口大小 2x2
        self.pool = nn.MaxPool2d(2, 2)
        # 全连接层：输入 64x8x8（展平），输出 512
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        # 全连接层：输入 512，输出 10（类别数）
        self.fc2 = nn.Linear(512, 10)
        # 激活函数
        self.relu = nn.ReLU()
        # Dropout 防止过拟合
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = self.relu(self.conv1(x))  # 卷积1 + ReLU
        x = self.pool(self.relu(self.conv2(x)))  # 卷积2 + ReLU + 池化
        x = x.view(-1, 64 * 8 * 8)  # 展平
        x = self.relu(self.fc1(x))  # 全连接1 + ReLU
        x = self.dropout(x)  # Dropout
        x = self.fc2(x)  # 全连接2
        return x

# 创建模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)

5. 定义损失函数与优化器

# 损失函数：交叉熵损失
criterion = nn.CrossEntropyLoss()

# 优化器：Adam
optimizer = optim.Adam(model.parameters(), lr=0.001)

6. 模型训练

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):  # 遍历数据集多个 epoch
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        # 梯度清零
        optimizer.zero_grad()

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()
        optimizer.step()

        # 打印损失
        running_loss += loss.item()
        if i % 100 == 99:  # 每 100 个 batch 打印一次
            print(f'[Epoch {epoch + 1}, Batch {i + 1}] loss: {running_loss / 100:.3f}')
            running_loss = 0.0

print('Finished Training')

7. 模型测试

# 测试模型
correct = 0
total = 0
model.eval()  # 设置模型为评估模式
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)  # 获取预测类别
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')

8. 按类别显示模型准确率

# 每个类别的准确率
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(len(labels)):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1

for i in range(10):
    print(f'Accuracy of {classes[i]} : {100 * class_correct[i] / class_total[i]:.2f}%')

9. 可视化模型预测

# 显示预测结果
dataiter = iter(testloader)
images, labels = next(dataiter)
outputs = model(images.to(device))
_, predicted = torch.max(outputs, 1)

# 显示前4张图片及其预测结果
imshow(torchvision.utils.make_grid(images[:4]))
print('Predicted: ', ' '.join(f'{classes[predicted[j]]}' for j in range(4)))

10. 总结

• 架构：构建了一个简单的 CNN，用于 CIFAR-10 数据集分类。
• 关键点：
  1. 数据加载与预处理：标准化和数据增强。
  2. 模型设计：卷积层 + 池化层 + 全连接层。
  3. 训练与测试：通过梯度下降更新参数。
• 改进方向：可以引入更复杂的网络（如 ResNet、VGG），或通过数据增强和正则化提升性能。

运行完上述代码后，CNN 的分类准确率通常可以达到 70% 左右。如果需要更高的准确率，可以调整超参数或使用更复杂的模型架构。