案例代码实现

一、代码说明

本案例使用PyTorch实现一个改进版LeNet-5模型，用于CIFAR-10数据集的图像分类任务。代码包含以下核心步骤：

数据加载与预处理（含数据增强，划分训练/验证/测试集）；

定义CNN网络结构（LeNet-5改进版，适配3通道输入）；

模型训练（用验证集评估泛化能力）；

模型测试与结果可视化（用独立测试集最终评估）。

适合人群：数模小白（无需深度学习基础，代码注释详细，逻辑清晰）。运行环境：Python3.8+、PyTorch1.10+、torchvision0.11+、matplotlib3.5+。

二、完整代码实现

# 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np# ------------------------------
# 1. 配置全局参数（数模小白可调整这里）
# ------------------------------
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 优先用GPU
BATCH_SIZE = 64  # 每批数据量（越大训练越快，但占内存越多）
EPOCHS = 10  # 训练轮数（越大模型越准，但训练时间越长）
LEARNING_RATE = 0.001  # 学习率（越小收敛越稳，但训练越慢）
VAL_SPLIT = 0.2  # 验证集占训练集的比例（20%）# ------------------------------
# 2. 数据加载与预处理（含数据增强，划分训练/验证/测试集）
# ------------------------------
def load_data():"""加载CIFAR-10数据集，返回训练/验证/测试DataLoader"""# 训练集数据增强（防止过拟合）：随机裁剪、水平翻转、归一化train_transform = transforms.Compose([transforms.RandomCrop(32, padding=4),  # 随机裁剪32x32，边缘补4像素transforms.RandomHorizontalFlip(),     # 随机水平翻转（50%概率）transforms.ToTensor(),                 # 转为Tensor（0-1）transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化到[-1,1]])# 验证集/测试集预处理（不增强，保持真实分布）val_test_transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])# 下载/加载数据集（第一次运行会下载，约170MB）full_train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=train_transform)val_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=val_test_transform)test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=val_test_transform)# 划分训练集和验证集（8:2）train_size = int((1 - VAL_SPLIT) * len(full_train_dataset))val_size = len(full_train_dataset) - train_sizetrain_dataset, _ = random_split(full_train_dataset, [train_size, val_size])_, val_dataset = random_split(val_dataset, [train_size, val_size])  # 保持验证集transform正确# 生成DataLoader（批量加载数据）train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)return train_loader, val_loader, test_loader# ------------------------------
# 3. 定义CNN网络结构（改进版LeNet-5）
# ------------------------------
class LeNet5(nn.Module):"""改进版LeNet-5，适配CIFAR-10的3通道输入（3x32x32）"""def __init__(self):super(LeNet5, self).__init__()# 卷积层1：提取边缘特征（3通道→6通道，5x5 kernel）self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5)# 最大池化层1：简化特征（2x2窗口，步长2）self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)# 卷积层2：提取纹理/形状特征（6通道→16通道，5x5 kernel）self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)# 最大池化层2：进一步简化特征（2x2窗口，步长2）self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)# 全连接层1：整合高级特征（16*5*5→120）self.fc1 = nn.Linear(16 * 5 * 5, 120)# 全连接层2：进一步整合特征（120→84）self.fc2 = nn.Linear(120, 84)# 输出层：分类决策（84→10类，对应CIFAR-10标签）self.fc3 = nn.Linear(84, 10)# 激活函数（ReLU，引入非线性，解决线性模型表达能力不足问题）self.relu = nn.ReLU()def forward(self, x):"""前向传播：定义数据在网络中的流动路径"""# 卷积层1 → ReLU → 池化层1：3x32x32 → 6x28x28 → 6x14x14x = self.pool1(self.relu(self.conv1(x)))# 卷积层2 → ReLU → 池化层2：6x14x14 → 16x10x10 → 16x5x5x = self.pool2(self.relu(self.conv2(x)))# 展平：将二维特征图转为一维向量（16x5x5 → 400），适配全连接层x = x.view(-1, 16 * 5 * 5)# 全连接层1 → ReLU：400 → 120x = self.relu(self.fc1(x))# 全连接层2 → ReLU：120 → 84x = self.relu(self.fc2(x))# 输出层：84 → 10（不使用Softmax，因为CrossEntropyLoss会自动处理）x = self.fc3(x)return x# ------------------------------
# 4. 模型训练与验证函数（用验证集评估泛化能力）
# ------------------------------
def train_model(model, train_loader, val_loader, optimizer, criterion):"""训练模型，每轮输出训练/验证损失与准确率"""best_val_acc = 0.0  # 记录最佳验证准确率（用于保存最优模型）for epoch in range(EPOCHS):# ------------------------------# 训练阶段（更新模型参数）# ------------------------------model.train()  # 切换到训练模式（启用BatchNorm/ Dropout等训练专用层）train_loss = 0.0train_correct = 0for inputs, labels in train_loader:inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)  # 数据移至GPU/CPUoptimizer.zero_grad()  # 清空梯度（避免梯度累积）outputs = model(inputs)  # 前向传播：输入→模型→输出（预测值）loss = criterion(outputs, labels)  # 计算损失（预测值与真实值的差距）loss.backward()  # 反向传播：计算梯度（从损失到各层参数）optimizer.step()  # 更新参数（用梯度调整参数，最小化损失）# 统计训练损失与准确率train_loss += loss.item() * inputs.size(0)  # 累计损失（乘以批量大小，避免批量大小影响）_, preds = torch.max(outputs, 1)  # 取预测概率最大的类别（0-9）train_correct += (preds == labels).sum().item()  # 统计正确预测的样本数# 计算训练集平均损失与准确率train_loss = train_loss / len(train_loader.dataset)train_acc = train_correct / len(train_loader.dataset)# ------------------------------# 验证阶段（评估泛化能力，不更新参数）# ------------------------------model.eval()  # 切换到验证模式（关闭BatchNorm/ Dropout等）val_loss = 0.0val_correct = 0with torch.no_grad():  # 关闭梯度计算（节省内存，加速验证）for inputs, labels in val_loader:inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)outputs = model(inputs)loss = criterion(outputs, labels)# 统计验证损失与准确率val_loss += loss.item() * inputs.size(0)_, preds = torch.max(outputs, 1)val_correct += (preds == labels).sum().item()# 计算验证集平均损失与准确率val_loss = val_loss / len(val_loader.dataset)val_acc = val_correct / len(val_loader.dataset)# 打印本轮训练/验证结果print(f"Epoch {epoch+1}/{EPOCHS}")print(f"训练集：损失={train_loss:.4f}，准确率={train_acc:.4f}")print(f"验证集：损失={val_loss:.4f}，准确率={val_acc:.4f}")print("-" * 50)# 保存最佳模型（验证准确率最高的模型，避免过拟合）if val_acc > best_val_acc:best_val_acc = val_acctorch.save(model.state_dict(), "best_model.pth")print(f"训练结束，最佳验证准确率={best_val_acc:.4f}（模型已保存至best_model.pth）")# ------------------------------
# 5. 模型测试与结果可视化（用独立测试集最终评估）
# ------------------------------
def test_model(model, test_loader):"""用独立测试集评估模型性能，输出准确率并可视化预测结果"""model.eval()  # 切换到验证模式test_correct = 0with torch.no_grad():  # 关闭梯度计算for inputs, labels in test_loader:inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)outputs = model(inputs)_, preds = torch.max(outputs, 1)test_correct += (preds == labels).sum().item()# 计算测试集准确率test_acc = test_correct / len(test_loader.dataset)print(f"\n测试集最终准确率={test_acc:.4f}")# 可视化10张测试图像的预测结果（直观展示模型效果）class_names = ["飞机", "汽车", "鸟", "猫", "鹿", "狗", "青蛙", "马", "船", "卡车"]inputs, labels = next(iter(test_loader))  # 取一批测试数据（BATCH_SIZE=64）inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)outputs = model(inputs)_, preds = torch.max(outputs, 1)# 绘制图像（2行5列，显示10张）plt.figure(figsize=(12, 6))for i in range(10):plt.subplot(2, 5, i+1)# 反归一化：将[-1,1]转回[0,1]（方便显示图像）img = inputs[i].cpu().numpy().transpose((1, 2, 0))  # 转为HWC格式（高度×宽度×通道）img = img * 0.5 + 0.5  # 反归一化（原归一化公式：img = (img - mean) / std → 反推：img = img * std + mean）plt.imshow(img)# 设置标题：真实标签 vs 预测标签plt.title(f"真实：{class_names[labels[i]]}\n预测：{class_names[preds[i]]}", fontsize=10)plt.axis("off")  # 隐藏坐标轴plt.tight_layout()  # 调整子图间距plt.show()# ------------------------------
# 6. 主程序（整合所有步骤，执行训练与测试）
# ------------------------------
if __name__ == "__main__":# 1. 加载数据（划分训练/验证/测试集）print("正在加载数据...")train_loader, val_loader, test_loader = load_data()print(f"数据加载完成：\n- 训练集大小：{len(train_loader.dataset)} \n- 验证集大小：{len(val_loader.dataset)} \n- 测试集大小：{len(test_loader.dataset)}")# 2. 初始化模型、损失函数、优化器print("\n正在初始化模型...")model = LeNet5().to(DEVICE)  # 将模型移至GPU/CPUcriterion = nn.CrossEntropyLoss()  # 交叉熵损失（适用于多分类任务）optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)  # Adam优化器（自适应学习率，收敛更稳定）# 3. 训练模型（用验证集评估）print("\n正在训练模型...")train_model(model, train_loader, val_loader, optimizer, criterion)# 4. 加载最佳模型并测试（用独立测试集）print("\n正在测试最佳模型...")model.load_state_dict(torch.load("best_model.pth"))  # 加载训练过程中保存的最佳模型test_model(model, test_loader)

三、代码使用说明

1.环境安装

打开命令行，运行以下命令安装依赖库（建议使用虚拟环境）：

pip install torch torchvision matplotlib numpy

2.运行代码

将代码保存为cnn_cifar10.py，在命令行中运行：

python cnn_cifar10.py

3.结果解释

训练过程：每轮（Epoch）输出训练集（更新参数）和验证集（评估泛化能力）的损失（Loss，越小说明预测越准）和准确率（Accuracy，越大说明模型越准）。

最佳模型：训练结束后，保存验证准确率最高的模型到best_model.pth（避免过拟合）。

测试结果：加载最佳模型后，用独立测试集评估，输出测试集准确率（一般在70%-85%之间，增加EPOCHS可提高），并显示10张测试图像的真实标签与预测标签（直观看到模型效果）。

四、数模小白调整建议

提高准确率：若训练集准确率低（<80%），可增加EPOCHS（如改为20），让模型多学习几轮；或增大LEARNING_RATE（如改为0.002），加快收敛速度。

缓解过拟合：若验证集准确率远低于训练集（如差 10% 以上），可添加更多数据增强（如transforms.RandomRotation(10)随机旋转 10 度、transforms.ColorJitter(brightness=0.2)调整亮度），或减小模型复杂度（如将conv1的out_channels=6改为3）。

加速训练：若训练太慢，可增大BATCH_SIZE（如改为128，需确保GPU内存足够），或使用更高效的优化器（如optim.AdamW，带权重衰减的Adam）。

五、常见问题解答

Q：为什么要划分验证集？A：验证集用于在训练过程中评估模型的泛化能力，避免模型“记住”训练集细节（过拟合）。测试集是最终评估模型性能的“考题”，不能在训练过程中使用。

Q：数据增强为什么有效？A：数据增强（如随机裁剪、翻转）通过生成“虚拟”训练数据，扩大了训练集的多样性，让模型学习到更通用的特征，从而提高泛化能力。

Q：为什么用Adam优化器而不是SGD？A：Adam优化器会为每个参数自适应调整学习率，比传统SGD（随机梯度下降）收敛更快、更稳定，适合新手使用。

通过运行这份代码，你可以完整体验CNN从数据预处理到模型部署的全流程，理解“卷积层提取特征、池化层简化特征、全连接层做决策”的核心逻辑，为后续更复杂的深度学习模型（如ResNet、YOLO）打下基础！