4. 卷积神经网络(CNN)基础

任务说明

本案例将使用PyTorch构建一个简单的卷积神经网络(CNN),用于CIFAR-10数据集的图像分类。CIFAR-10包含10个类别的32×32彩色图像(飞机、汽车、鸟、猫、鹿、狗、青蛙、马、船、卡车),训练集50000张,测试集10000张。目标是通过训练CNN模型,在测试集上达到可验证的分类准确率。

模型结构

模型采用经典的卷积-池化-全连接结构,包含两个卷积块(每个卷积块由卷积层、ReLU激活和最大池化组成),随后将特征图展平并送入两个全连接层进行分类。输入为3通道32×32图像,输出为10类概率分布。

输入 3×32×32 Conv1 3→16通道 3×3, ReLU Pool1 2×2最大池化 Conv2 16→32通道 3×3, ReLU Pool2 2×2最大池化 展平 8×8×32 全连接层 FC1: 2048→256 FC2: 256→10 输出 10类概率

代码实现


import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 1. 定义CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv_block1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.conv_block2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        x = self.conv_block1(x)
        x = self.conv_block2(x)
        x = self.fc(x)
        return x

# 2. 数据加载与预处理
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, num_workers=2)

testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=64, shuffle=False, num_workers=2)

# 3. 初始化模型、损失函数与优化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. 训练一个epoch
def train_one_epoch():
    model.train()
    running_loss = 0.0
    for images, labels in trainloader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    return running_loss / len(trainloader)

# 5. 测试准确率
def evaluate():
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in testloader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100.0 * correct / total

# 6. 训练与验证
if __name__ == '__main__':
    epochs = 5
    for epoch in range(epochs):
        loss = train_one_epoch()
        acc = evaluate()
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss:.4f}, Test Acc: {acc:.2f}%')

运行方式

确保已安装torch和torchvision,执行以下命令运行脚本(首次运行会自动下载CIFAR-10数据集,约170MB):


python cnn_cifar10.py

验证方法

训练过程中每个epoch结束后会打印测试集准确率。判定标准:

  • 程序能正常完成5个epoch的训练,无报错;
  • 每个epoch的Loss值呈下降趋势;
  • 测试准确率随训练进行逐步提升,5个epoch后预期达到60%以上(随机猜测为10%,训练收敛后应显著高于此值)。

若使用GPU训练,每个epoch约需30秒;纯CPU训练约需3-5分钟。可通过调整epochs变量增加训练轮次以追求更高准确率。