基于PyTorch的图像分类实战：完整代码与深度解析

作者：十万个为什么2025.10.12 01:01浏览量：2

简介：本文提供基于PyTorch的图像分类完整实现，包含数据加载、模型构建、训练与评估全流程代码及详细注释，适合开发者快速掌握深度学习图像分类技术。

基于PyTorch的图像分类实战：完整代码与深度解析

引言

图像分类是计算机视觉领域的核心任务，PyTorch作为主流深度学习框架，凭借其动态计算图和Pythonic接口，成为开发者实现图像分类的首选工具。本文将通过一个完整的实战案例，展示如何使用PyTorch实现从数据加载到模型部署的全流程，并提供详细注释帮助读者理解每个步骤的实现原理。

一、环境准备与数据集选择

1.1 环境配置

首先需要安装PyTorch及相关依赖库：

  1. pip install torch torchvision matplotlib numpy

建议使用GPU环境加速训练，可通过torch.cuda.is_available()检查CUDA是否可用。

1.2 数据集选择

本文使用CIFAR-10数据集，该数据集包含10个类别的60000张32x32彩色图像，分为50000张训练集和10000张测试集。PyTorch的torchvision.datasets模块提供了便捷的数据加载接口：

  1. import torchvision
  2. import torchvision.transforms as transforms
  4. # 定义数据预处理流程
  5. transform = transforms.Compose([
  6.     transforms.ToTensor(),  # 将PIL图像或numpy数组转为Tensor，并缩放到[0,1]
  7.     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化到[-1,1]
  8. ])
  10. # 加载训练集和测试集
  11. trainset = torchvision.datasets.CIFAR10(
  12.     root='./data', train=True, download=True, transform=transform)
  13. trainloader = torch.utils.data.DataLoader(
  14.     trainset, batch_size=32, shuffle=True, num_workers=2)
  16. testset = torchvision.datasets.CIFAR10(
  17.     root='./data', train=False, download=True, transform=transform)
  18. testloader = torch.utils.data.DataLoader(
  19.     testset, batch_size=32, shuffle=False, num_workers=2)
  21. # 定义类别标签
  22. classes = ('plane', 'car', 'bird', 'cat', 'deer', 
  23.            'dog', 'frog', 'horse', 'ship', 'truck')

关键点解析：

  • transforms.Compose：将多个数据预处理操作组合成一个流水线
  • Normalize：使用均值和标准差进行标准化，这里采用(0.5,0.5,0.5)的简单缩放
  • DataLoader：实现批量加载、数据打乱和多线程加速

二、CNN模型构建

2.1 基础CNN模型实现

我们构建一个包含3个卷积层和2个全连接层的CNN模型：

  1. import torch.nn as nn
  2. import torch.nn.functional as F
  4. class CNN(nn.Module):
  5.     def __init__(self):
  6.         super(CNN, self).__init__()
  7.         # 卷积层1：输入通道3，输出通道32，3x3卷积核
  8.         self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
  9.         # 卷积层2：输入通道32，输出通道64，3x3卷积核
  10.         self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
  11.         # 卷积层3：输入通道64，输出通道128，3x3卷积核
  12.         self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
  14.         # 最大池化层
  15.         self.pool = nn.MaxPool2d(2, 2)
  17.         # 全连接层
  18.         self.fc1 = nn.Linear(128 * 4 * 4, 512)  # 经过3次池化后，特征图尺寸为4x4
  19.         self.fc2 = nn.Linear(512, 10)  # 输出10个类别
  21.     def forward(self, x):
  22.         # 第一层卷积+池化+ReLU
  23.         x = self.pool(F.relu(self.conv1(x)))
  24.         # 第二层卷积+池化+ReLU
  25.         x = self.pool(F.relu(self.conv2(x)))
  26.         # 第三层卷积+池化+ReLU
  27.         x = self.pool(F.relu(self.conv3(x)))
  29.         # 展平特征图
  30.         x = x.view(-1, 128 * 4 * 4)
  32.         # 全连接层
  33.         x = F.relu(self.fc1(x))
  34.         x = self.fc2(x)
  35.         return x

模型结构解析：

  1. 输入层：32x32x3的RGB图像
  2. 卷积块1：32个3x3卷积核 → ReLU激活 → 2x2最大池化 → 输出尺寸16x16x32
  3. 卷积块2：64个3x3卷积核 → ReLU激活 → 2x2最大池化 → 输出尺寸8x8x64
  4. 卷积块3：128个3x3卷积核 → ReLU激活 → 2x2最大池化 → 输出尺寸4x4x128
  5. 全连接层：展平为2048维向量 → 512维隐藏层 → 10维输出

2.2 模型初始化与设备分配

  1. import torch
  3. # 初始化模型
  4. model = CNN()
  5. # 将模型移动到GPU（如果可用）
  6. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  7. model.to(device)

三、训练流程实现

3.1 定义损失函数和优化器

  1. import torch.optim as optim
  3. # 交叉熵损失函数
  4. criterion = nn.CrossEntropyLoss()
  5. # 随机梯度下降优化器，学习率0.001，动量0.9
  6. optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

3.2 训练循环实现

  1. def train_model(model, trainloader, criterion, optimizer, epochs=10):
  2.     for epoch in range(epochs):
  3.         running_loss = 0.0
  4.         correct = 0
  5.         total = 0
  7.         for i, data in enumerate(trainloader, 0):
  8.             # 获取输入和标签
  9.             inputs, labels = data[0].to(device), data[1].to(device)
  11.             # 梯度清零
  12.             optimizer.zero_grad()
  14.             # 前向传播
  15.             outputs = model(inputs)
  17.             # 计算损失
  18.             loss = criterion(outputs, labels)
  20.             # 反向传播和优化
  21.             loss.backward()
  22.             optimizer.step()
  24.             # 统计信息
  25.             running_loss += loss.item()
  26.             _, predicted = torch.max(outputs.data, 1)
  27.             total += labels.size(0)
  28.             correct += (predicted == labels).sum().item()
  30.             # 每200个batch打印一次统计信息
  31.             if i % 200 == 199:
  32.                 print(f'Epoch {epoch+1}, Batch {i+1}, '
  33.                       f'Loss: {running_loss/200:.3f}, '
  34.                       f'Acc: {100*correct/total:.2f}%')
  35.                 running_loss = 0.0
  37.         # 每个epoch结束后打印验证准确率
  38.         test_acc = evaluate_model(model, testloader)
  39.         print(f'Epoch {epoch+1} completed, Test Acc: {test_acc:.2f}%')
  41. def evaluate_model(model, testloader):
  42.     correct = 0
  43.     total = 0
  44.     with torch.no_grad():
  45.         for data in testloader:
  46.             images, labels = data[0].to(device), data[1].to(device)
  47.             outputs = model(images)
  48.             _, predicted = torch.max(outputs.data, 1)
  49.             total += labels.size(0)
  50.             correct += (predicted == labels).sum().item()
  51.     return 100 * correct / total

训练流程解析：

  1. 外层循环控制训练轮次（epochs）
  2. 内层循环遍历每个batch的数据
  3. 关键步骤：梯度清零→前向传播→计算损失→反向传播→参数更新
  4. 统计训练损失和准确率
  5. 每个epoch结束后在测试集上评估模型性能

四、完整代码与运行示例

4.1 完整实现代码

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4. import torch.optim as optim
  5. import torchvision
  6. import torchvision.transforms as transforms
  7. import matplotlib.pyplot as plt
  8. import numpy as np
  10. # 设备配置
  11. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  13. # 1. 数据加载
  14. transform = transforms.Compose([
  15.     transforms.ToTensor(),
  16.     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  17. ])
  19. trainset = torchvision.datasets.CIFAR10(
  20.     root='./data', train=True, download=True, transform=transform)
  21. trainloader = torch.utils.data.DataLoader(
  22.     trainset, batch_size=32, shuffle=True, num_workers=2)
  24. testset = torchvision.datasets.CIFAR10(
  25.     root='./data', train=False, download=True, transform=transform)
  26. testloader = torch.utils.data.DataLoader(
  27.     testset, batch_size=32, shuffle=False, num_workers=2)
  29. classes = ('plane', 'car', 'bird', 'cat', 'deer', 
  30.            'dog', 'frog', 'horse', 'ship', 'truck')
  32. # 2. 模型定义
  33. class CNN(nn.Module):
  34.     def __init__(self):
  35.         super(CNN, self).__init__()
  36.         self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
  37.         self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
  38.         self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
  39.         self.pool = nn.MaxPool2d(2, 2)
  40.         self.fc1 = nn.Linear(128 * 4 * 4, 512)
  41.         self.fc2 = nn.Linear(512, 10)
  43.     def forward(self, x):
  44.         x = self.pool(F.relu(self.conv1(x)))
  45.         x = self.pool(F.relu(self.conv2(x)))
  46.         x = self.pool(F.relu(self.conv3(x)))
  47.         x = x.view(-1, 128 * 4 * 4)
  48.         x = F.relu(self.fc1(x))
  49.         x = self.fc2(x)
  50.         return x
  52. model = CNN().to(device)
  54. # 3. 训练配置
  55. criterion = nn.CrossEntropyLoss()
  56. optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
  58. # 4. 训练函数
  59. def train_model(model, trainloader, criterion, optimizer, epochs=10):
  60.     for epoch in range(epochs):
  61.         running_loss = 0.0
  62.         correct = 0
  63.         total = 0
  65.         for i, data in enumerate(trainloader, 0):
  66.             inputs, labels = data[0].to(device), data[1].to(device)
  68.             optimizer.zero_grad()
  69.             outputs = model(inputs)
  70.             loss = criterion(outputs, labels)
  71.             loss.backward()
  72.             optimizer.step()
  74.             running_loss += loss.item()
  75.             _, predicted = torch.max(outputs.data, 1)
  76.             total += labels.size(0)
  77.             correct += (predicted == labels).sum().item()
  79.             if i % 200 == 199:
  80.                 print(f'Epoch {epoch+1}, Batch {i+1}, '
  81.                       f'Loss: {running_loss/200:.3f}, '
  82.                       f'Acc: {100*correct/total:.2f}%')
  83.                 running_loss = 0.0
  85.         test_acc = evaluate_model(model, testloader)
  86.         print(f'Epoch {epoch+1} completed, Test Acc: {test_acc:.2f}%')
  88. def evaluate_model(model, testloader):
  89.     correct = 0
  90.     total = 0
  91.     with torch.no_grad():
  92.         for data in testloader:
  93.             images, labels = data[0].to(device), data[1].to(device)
  94.             outputs = model(images)
  95.             _, predicted = torch.max(outputs.data, 1)
  96.             total += labels.size(0)
  97.             correct += (predicted == labels).sum().item()
  98.     return 100 * correct / total
  100. # 5. 执行训练
  101. train_model(model, trainloader, criterion, optimizer, epochs=10)
  103. # 6. 保存模型
  104. torch.save(model.state_dict(), 'cifar_cnn.pth')

4.2 运行结果与优化建议

典型运行结果：

  1. Epoch 1, Batch 200, Loss: 2.034, Acc: 32.45%
  2. Epoch 1, Batch 400, Loss: 1.876, Acc: 36.78%
  3. ...
  4. Epoch 10 completed, Test Acc: 72.34%

优化建议：

  1. 学习率调整：使用学习率调度器（如torch.optim.lr_scheduler.StepLR）
  2. 数据增强：添加随机裁剪、水平翻转等增强操作
  3. 模型改进：尝试ResNet等更先进的架构
  4. 超参数调优：使用网格搜索或贝叶斯优化

五、总结与扩展

本文完整实现了基于PyTorch的图像分类系统，涵盖了从数据加载到模型部署的全流程。关键技术点包括：

  1. 使用torchvision高效加载和预处理数据
  2. 构建包含卷积层、池化层和全连接层的CNN模型
  3. 实现完整的训练循环和评估流程
  4. 提供GPU加速支持和模型保存功能

扩展方向：

  • 尝试迁移学习，使用预训练模型（如ResNet）进行微调
  • 实现更复杂的数据增强策略
  • 部署模型到移动端或Web服务
  • 探索更先进的模型架构（如EfficientNet、Vision Transformer）

通过本文的实践，读者可以掌握PyTorch实现图像分类的核心技术，为后续更复杂的计算机视觉任务打下坚实基础。

最热文章