PyTorch 实战 —— 从零搭建 CV 模型（图像分类完整流程）
前言：为什么选择 “图像分类” 作为入门任务？
计算机视觉（CV）的核心任务包括图像分类、目标检测、语义分割等，其中图像分类是最基础、最易上手的任务 —— 只需判断 “一张图片属于哪一类”（如猫 / 狗、汽车 / 飞机），适合初学者理解 CV 模型的核心逻辑（特征提取→分类决策）。
本文将基于 PyTorch 框架，从零搭建一个卷积神经网络（CNN） 用于 CIFAR-10 数据集分类（10 个类别：飞机、汽车、鸟类、猫、鹿、狗、青蛙、马、船、卡车），全程覆盖 “环境准备→数据预处理→模型搭建→训练→评估→预测”，代码可直接复制运行，新手也能快速上手。
一、环境准备：PyTorch 及依赖库安装
首先确保你的环境已安装 PyTorch 及 CV 相关工具库，推荐使用 Anaconda 管理环境（避免版本冲突）。
1. 创建并激活虚拟环境（可选但推荐）
|
# 创建名为 pytorch-cv 的虚拟环境（Python 3.9 兼容性较好） conda create -n pytorch-cv python=3.9 # 激活环境（Windows） conda activate pytorch-cv # 激活环境（Linux/macOS） source activate pytorch-cv |
2. 安装核心库
torch/torchvision：PyTorch 核心框架及 CV 工具集（含数据集、模型组件）；matplotlib：用于绘制训练曲线、显示图片；numpy：基础数值计算（PyTorch 已集成大部分功能，仅备用）。
|
# 安装 PyTorch（根据你的 CUDA 版本选择，无 GPU 选 CPU 版本） # 1. 无 GPU（CPU 版本） conda install pytorch torchvision torchaudio cpuonly -c pytorch # 2. 有 GPU（CUDA 11.8 示例，需先安装对应版本的 CUDA Toolkit） conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia # 安装其他依赖 conda install matplotlib numpy |
3. 验证环境是否成功
运行以下代码，若能正常输出 PyTorch 版本和 GPU 状态（有 GPU 会显示 True），则环境没问题：
|
import torch import torchvision # 查看 PyTorch 版本 print(“PyTorch 版本：”, torch.__version__) # 查看 GPU 是否可用 print(“GPU 是否可用：”, torch.cuda.is_available()) # 查看 GPU 设备数量（有 GPU 会显示 ≥1） print(“GPU 设备数量：”, torch.cuda.device_count()) # 查看当前使用的 GPU（有 GPU 会显示设备号，如 0） if torch.cuda.is_available(): print(“当前使用的 GPU：”, torch.cuda.get_device_name(0)) |
二、数据集加载与预处理：CIFAR-10 实战
PyTorch 的 torchvision.datasets 模块已内置 CIFAR-10 数据集，无需手动下载，直接调用即可。但 raw 数据无法直接输入模型，需做数据预处理（标准化、增强等）。
1. 数据预处理逻辑
数据增强：仅对训练集使用（避免过拟合），如随机裁剪、水平翻转，模拟不同场景下的图片；标准化：对训练集和测试集均使用，将像素值从 [0,255] 归一化到 [-1,1] 左右（基于 CIFAR-10 数据集的均值和标准差），让模型更容易收敛。
2. 代码实现：加载并预处理数据
|
import torchvision.transforms as transforms from torchvision.datasets import CIFAR10 from torch.utils.data import DataLoader # 1. 定义数据预处理管道 # 训练集：数据增强 + 标准化 train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), # 随机裁剪（原始32×32， padding=4后先扩为40×40再裁回32×32） transforms.RandomHorizontalFlip(p=0.5), # 50%概率水平翻转 transforms.ToTensor(), # 转为 Tensor（像素值从 [0,255] 转为 [0,1]） transforms.Normalize( # 标准化：mean和std是CIFAR-10数据集的全局均值和标准差 mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616] ) ]) # 测试集：仅标准化（不做增强，确保评估真实性能） test_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize( mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616] ) ]) # 2. 加载 CIFAR-10 数据集 # 训练集（root：数据保存路径，train=True 表示训练集，download=True 自动下载） train_dataset = CIFAR10( root=”./data”, train=True, download=True, transform=train_transform ) # 测试集（train=False 表示测试集） test_dataset = CIFAR10( root=”./data”, train=False, download=True, transform=test_transform ) # 3. 创建 DataLoader（批量加载数据，支持多线程和打乱） batch_size = 64 # 每次输入模型的图片数量（根据GPU内存调整，小内存可设32） train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, # 训练集打乱，增强随机性 num_workers=2 # 多线程加载数据（Windows建议设0，避免线程报错） ) test_loader = DataLoader( test_dataset, batch_size=batch_size, shuffle=False, # 测试集无需打乱 num_workers=2 ) # 查看数据集基本信息 print(“训练集样本数：”, len(train_dataset)) # 输出 50000（CIFAR-10 训练集共5万张） print(“测试集样本数：”, len(test_dataset)) # 输出 10000（测试集共1万张） print(“类别列表：”, train_dataset.classes) # 输出 10个类别的名称 |
三、模型搭建：从零实现基础 CNN 网络
卷积神经网络（CNN）的核心是 “卷积层提取特征 + 全连接层分类”，我们搭建一个轻量型 CNN，结构如下（适合 CIFAR-10 小图片）：
输入（32x32x3）→ 卷积层1 → 激活函数 → 池化层1 → 卷积层2 → 激活函数 → 池化层2 → 全连接层1 → 激活函数 → 全连接层2（输出10类）
1. 代码实现：定义 CNN 模型
|
import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self, num_classes=10): super(SimpleCNN, self).__init__() # 1. 卷积层1：输入3通道（RGB），输出32通道，卷积核3×3，步长1， padding=1（保证输出尺寸不变） self.conv1 = nn.Conv2d( in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1 ) # 池化层1：最大池化，核2×2，步长2（输出尺寸减半：32×32 → 16×16） self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
# 2. 卷积层2：输入32通道，输出64通道，卷积核3×3，步长1， padding=1 self.conv2 = nn.Conv2d( in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1 ) # 池化层2：最大池化，核2×2，步长2（输出尺寸减半：16×16 → 8×8） self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
# 3. 全连接层1：输入特征数 = 64通道 × 8×8尺寸（池化后特征图大小），输出512维 self.fc1 = nn.Linear(64 * 8 * 8, 512) # 全连接层2：输入512维，输出10维（对应10个类别） self.fc2 = nn.Linear(512, num_classes) # 前向传播（模型的核心逻辑，定义数据如何流过网络） def forward(self, x): # 卷积1 → ReLU激活 → 池化1 x = self.pool1(F.relu(self.conv1(x))) # 卷积2 → ReLU激活 → 池化2 x = self.pool2(F.relu(self.conv2(x))) # 展平特征图：从 (batch_size, 64, 8, 8) 转为 (batch_size, 64*8*8) x = x.view(-1, 64 * 8 * 8) # -1 表示自动计算 batch_size # 全连接1 → ReLU激活 x = F.relu(self.fc1(x)) # 全连接2（输出 logits，后续用 softmax 转概率） x = self.fc2(x) return x # 实例化模型 model = SimpleCNN(num_classes=10) # 打印模型结构（查看是否符合预期） print(model) # 将模型移动到 GPU（若有 GPU） device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”) model.to(device) print(f”模型运行设备：{device}”) |
四、模型训练：定义损失、优化器与训练循环
训练的核心是 “最小化损失函数”—— 通过优化器调整模型参数，让模型的预测结果越来越接近真实标签。
1. 定义核心组件
损失函数：用 CrossEntropyLoss（交叉熵损失），适合多分类任务，内置了 Softmax 函数（无需手动在模型输出后加 Softmax）；优化器：用 Adam（自适应学习率优化器），比传统的 SGD 收敛更快，适合新手；训练轮次（epoch）：整个训练集遍历一次为 1 个 epoch，这里设 20 个 epoch（足够看到收敛趋势）。
2. 代码实现：训练循环
|
import torch.optim as optim import time import matplotlib.pyplot as plt # 1. 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() # 交叉熵损失（多分类） optimizer = optim.Adam(model.parameters(), lr=1e-3) # Adam优化器，学习率1e-3 # 2. 训练参数 num_epochs = 20 # 训练轮次 train_loss_history = [] # 记录训练损失（用于后续绘图） train_acc_history = [] # 记录训练准确率 test_acc_history = [] # 记录测试准确率 # 3. 训练循环 start_time = time.time() # 计时，查看训练耗时 for epoch in range(num_epochs): # ———————- 训练阶段 ———————- model.train() # 切换到训练模式（启用 Dropout、BatchNorm 等训练特有的层） running_loss = 0.0 # 记录当前 epoch 的总损失 correct = 0 # 记录当前 epoch 训练集正确预测的样本数 total = 0 # 记录当前 epoch 训练集总样本数 # 遍历训练集 DataLoader（每次取一个 batch） for i, (images, labels) in enumerate(train_loader): # 将数据移动到 GPU（若有） images, labels = images.to(device), labels.to(device) # 前向传播：计算模型预测 outputs = model(images) # 计算损失 loss = criterion(outputs, labels) # 反向传播 + 优化器更新参数 optimizer.zero_grad() # 清空上一轮的梯度（避免累积） loss.backward() # 反向传播计算梯度 optimizer.step() # 优化器更新模型参数 # 统计损失和准确率 running_loss += loss.item() * images.size(0) # 累计损失（乘以 batch_size，避免受 batch 大小影响） _, predicted = torch.max(outputs.data, 1) # 取预测概率最大的类别（outputs.data 是 logits，torch.max 取索引） total += labels.size(0) correct += (predicted == labels).sum().item() # 统计正确预测数 # 计算当前 epoch 的平均损失和准确率 epoch_train_loss = running_loss / len(train_loader.dataset) epoch_train_acc = correct / total * 100 train_loss_history.append(epoch_train_loss) train_acc_history.append(epoch_train_acc) # ———————- 评估阶段 ———————- model.eval() # 切换到评估模式（禁用 Dropout、固定 BatchNorm 统计量） test_correct = 0 test_total = 0 # 评估时不计算梯度（节省内存，加速计算） with torch.no_grad(): for images, labels in test_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) test_total += labels.size(0) test_correct += (predicted == labels).sum().item() epoch_test_acc = test_correct / test_total * 100 test_acc_history.append(epoch_test_acc) # 打印当前 epoch 的训练结果 print(f”Epoch [{epoch+1}/{num_epochs}], “ f”Train Loss: {epoch_train_loss:.4f}, “ f”Train Acc: {epoch_train_acc:.2f}%, “ f”Test Acc: {epoch_test_acc:.2f}%, “ f”Time Elapsed: {time.time()-start_time:.2f}s”) # 训练结束，打印总耗时 print(f”Training Finished! Total Time: {time.time()-start_time:.2f}s”) |
五、结果可视化与模型保存
训练完成后，我们需要通过可视化训练曲线分析模型性能（是否过拟合、是否收敛），并保存训练好的模型（方便后续预测）。
1. 可视化训练曲线
|
# 设置中文字体（避免 matplotlib 显示中文乱码） plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans'] plt.rcParams['axes.unicode_minus'] = False # 创建画布 fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) # 1. 绘制损失曲线 ax1.plot(range(1, num_epochs+1), train_loss_history, 'b-', label='训练损失') ax1.set_xlabel('训练轮次（Epoch）') ax1.set_ylabel('损失值') ax1.set_title('训练损失变化曲线') ax1.legend() ax1.grid(True) # 2. 绘制准确率曲线 ax2.plot(range(1, num_epochs+1), train_acc_history, 'r-', label='训练准确率') ax2.plot(range(1, num_epochs+1), test_acc_history, 'g-', label='测试准确率') ax2.set_xlabel('训练轮次（Epoch）') ax2.set_ylabel('准确率（%）') ax2.set_title('训练/测试准确率变化曲线') ax2.legend() ax2.grid(True) # 保存图片（可选） plt.savefig('train_curve.png', dpi=300, bbox_inches='tight') # 显示图片 plt.show() |
2. 保存与加载模型
PyTorch 保存模型有两种常用方式：
保存整个模型（含结构 + 参数）：简单，但不灵活（换环境可能因版本问题报错）；仅保存参数（推荐）：灵活，只需重新定义模型结构即可加载参数。
|
# 方式1：仅保存模型参数（推荐） torch.save(model.state_dict(), 'simple_cnn_cifar10.pth') print(“模型参数已保存到 simple_cnn_cifar10.pth”) # 方式2：保存整个模型（不推荐，兼容性差） # torch.save(model, 'simple_cnn_cifar10_full.pth') # ———————- 加载模型（后续预测时使用） ———————- # 1. 重新定义模型结构（需与训练时一致） loaded_model = SimpleCNN(num_classes=10) # 2. 加载保存的参数 loaded_model.load_state_dict(torch.load('simple_cnn_cifar10.pth')) # 3. 移动到设备并切换到评估模式 loaded_model.to(device) loaded_model.eval() print(“模型加载完成，可用于预测”) |
六、模型预测：用训练好的模型分类新图片
训练好的模型需要能对 “新图片” 进行分类，这里我们从测试集中随机选一张图片，演示预测流程。
1. 代码实现：单张图片预测
|
import random from PIL import Image # 1. 定义预测函数（输入图片路径，输出预测类别） def predict_image(image_path, model, transform, classes): # 加载图片（用 PIL 读取，保持与数据集一致的格式） image = Image.open(image_path).convert('RGB') # 转为 RGB（避免灰度图问题） # 预处理（与测试集预处理一致） image_tensor = transform(image).unsqueeze(0) # unsqueeze(0) 增加 batch 维度（模型需要 batch 输入） # 移动到设备 image_tensor = image_tensor.to(device)
# 预测 model.eval() with torch.no_grad(): outputs = model(image_tensor) _, predicted = torch.max(outputs.data, 1) # 取预测类别索引 predicted_class = classes[predicted.item()] # 转为类别名称 # 计算预测概率（可选） probabilities = F.softmax(outputs, dim=1) max_prob = probabilities[0][predicted.item()].item() * 100 # 最大概率
return image, predicted_class, max_prob # 2. 从测试集中随机选一张图片（模拟“新图片”） # 随机选一个测试集样本的索引 random_idx = random.randint(0, len(test_dataset)-1) # 获取图片和真实标签 test_image, test_label = test_dataset[random_idx] # 将 Tensor 转回 PIL 图片（用于显示） inv_transform = transforms.Compose([ transforms.Normalize(mean=[-0.4914/0.2470, -0.4822/0.2435, -0.4465/0.2616], # 逆标准化：恢复到 [0,1] std=[1/0.2470, 1/0.2435, 1/0.2616]), transforms.ToPILImage() ]) pil_image = inv_transform(test_image) # 保存图片到本地（作为“新图片”路径） image_path = 'test_image.png' pil_image.save(image_path) # 3. 调用预测函数 classes = train_dataset.classes # 类别列表 image, predicted_class, max_prob = predict_image( image_path=image_path, model=loaded_model, transform=test_transform, classes=classes ) # 4. 显示结果 plt.figure(figsize=(6, 6)) plt.imshow(image) plt.title(f”真实类别：{classes[test_label]} plt.axis('off') # 隐藏坐标轴 plt.show() |
七、进阶优化方向（可选）
本文搭建的是基础 CNN 模型，测试准确率约 70%-80%，若想进一步提升性能，可尝试以下优化方向：
模型改进：使用经典 CNN 结构（如 ResNet-18、MobileNetV2），解决梯度消失问题，提升特征提取能力；数据增强：增加更多增强手段（如 CutMix、MixUp、随机旋转），进一步提升泛化能力；正则化：添加 Dropout 层（如在全连接层前加 nn.Dropout(p=0.5)）、L2 正则化（优化器中加 weight_decay=1e-4），抑制过拟合；学习率调度：使用 ReduceLROnPlateau 或 StepLR 动态调整学习率（训练后期减小学习率，让模型更稳定收敛）；迁移学习：基于预训练模型（如 ImageNet 预训练的 ResNet）微调，适合小数据集或需快速提升性能的场景。
八、常见问题解决
训练时 GPU 内存不足：减小 batch_size（如从 64 改为 32）、使用更小的模型（如减少卷积层通道数）；训练损失不下降：检查学习率（太大可能震荡，太小收敛慢）、确认数据预处理是否正确（如均值 / 标准差是否匹配数据集）；测试准确率远低于训练准确率：存在过拟合，需增加数据增强、添加正则化层；Windows 系统 DataLoader 线程报错：将 num_workers 设为 0（禁用多线程）。
总结
本文从零搭建了一个基于 PyTorch 的 CV 图像分类模型，核心流程可总结为：
环境准备 → 数据加载与预处理 → 模型定义（CNN） → 训练（损失+优化器） → 评估与可视化 → 预测
通过这个实战，你不仅掌握了 PyTorch 搭建 CV 模型的基础能力，更理解了每个步骤的核心逻辑（如数据增强的目的、CNN 各层的作用）。后续可尝试将模型扩展到其他 CV 任务（如目标检测用 Faster R-CNN、语义分割用 U-Net），逐步深入计算机视觉领域。
暂无评论内容