logo

开发者热搜

从零开始：使用PyTorch实现图像分类（附完整代码与注释）

作者：狼烟四起2025.09.18 17:43浏览量：0

简介：本文详细介绍如何使用PyTorch框架实现图像分类任务，包含数据加载、模型构建、训练与评估全流程，提供完整可运行的代码及逐行注释，适合PyTorch初学者和深度学习实践者。

从零开始：使用PyTorch实现图像分类（附完整代码与注释）

引言

图像分类是计算机视觉领域的核心任务之一，广泛应用于人脸识别、医学影像分析、自动驾驶等场景。PyTorch作为主流深度学习框架，以其动态计算图和简洁的API设计，成为实现图像分类任务的理想选择。本文将通过一个完整的案例，展示如何使用PyTorch从零开始实现图像分类，涵盖数据准备、模型构建、训练优化和结果评估的全流程，并提供详细注释的完整代码。

1. 环境准备与依赖安装

在开始之前，需要确保已安装Python 3.6+环境，并安装以下依赖库：

  1. pip install torch torchvision matplotlib numpy

PyTorch提供预编译的二进制包，支持CPU和GPU版本。若使用GPU加速，需安装与CUDA版本匹配的PyTorch版本。

2. 数据准备与预处理

2.1 数据集选择

本文使用经典的CIFAR-10数据集，包含10个类别的6万张32x32彩色图像（5万训练集，1万测试集）。PyTorch的torchvision库提供了便捷的数据加载接口：

  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', 
  13.     train=True,
  14.     download=True, 
  15.     transform=transform
  16. )
  17. trainloader = torch.utils.data.DataLoader(
  18.     trainset, 
  19.     batch_size=32,
  20.     shuffle=True, 
  21.     num_workers=2
  22. )
  24. testset = torchvision.datasets.CIFAR10(
  25.     root='./data', 
  26.     train=False,
  27.     download=True, 
  28.     transform=transform
  29. )
  30. testloader = torch.utils.data.DataLoader(
  31.     testset, 
  32.     batch_size=32,
  33.     shuffle=False, 
  34.     num_workers=2
  35. )

关键点说明：

  • transforms.Compose：将多个预处理操作组合为一个流水线。
  • ToTensor()：自动将图像从[0,255]的uint8类型转换为[0,1]的float32类型Tensor。
  • Normalize：使用均值和标准差进行归一化，公式为(x - mean) / std。
  • DataLoader：提供批量加载、多线程加速和随机打乱功能。

2.2 数据可视化

为验证数据加载是否正确，可随机显示部分训练图像：

  1. import matplotlib.pyplot as plt
  2. import numpy as np
  4. def imshow(img):
  5.     img = img / 2 + 0.5  # 反归一化
  6.     npimg = img.numpy()
  7.     plt.imshow(np.transpose(npimg, (1, 2, 0)))
  8.     plt.show()
  10. # 获取一个批次的图像和标签
  11. dataiter = iter(trainloader)
  12. images, labels = next(dataiter)
  14. # 显示图像
  15. imshow(torchvision.utils.make_grid(images))
  16. # 打印标签
  17. print(' '.join(f'{trainset.classes[labels[j]]}' for j in range(4)))

3. 模型构建：卷积神经网络（CNN）

3.1 网络架构设计

本文实现一个简化的CNN模型，包含2个卷积层和2个全连接层：

  1. import torch.nn as nn
  2. import torch.nn.functional as F
  4. class Net(nn.Module):
  5.     def __init__(self):
  6.         super(Net, self).__init__()
  7.         # 第一个卷积层：输入通道3（RGB），输出通道6，5x5卷积核
  8.         self.conv1 = nn.Conv2d(3, 6, 5)
  9.         # 第二个卷积层：输入通道6，输出通道16，5x5卷积核
  10.         self.conv2 = nn.Conv2d(6, 16, 5)
  11.         # 第一个全连接层：输入16*5*5（经过池化后的特征图尺寸），输出120
  12.         self.fc1 = nn.Linear(16 * 5 * 5, 120)
  13.         # 第二个全连接层：输入120，输出84
  14.         self.fc2 = nn.Linear(120, 84)
  15.         # 输出层：输入84，输出10（CIFAR-10类别数）
  16.         self.fc3 = nn.Linear(84, 10)
  18.     def forward(self, x):
  19.         # 最大池化，kernel_size=2, stride=2
  20.         x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
  21.         x = F.max_pool2d(F.relu(self.conv2(x)), 2)
  22.         # 展平特征图
  23.         x = x.view(-1, 16 * 5 * 5)
  24.         x = F.relu(self.fc1(x))
  25.         x = F.relu(self.fc2(x))
  26.         x = self.fc3(x)
  27.         return x
  29. net = Net()

架构解析：

  • 卷积层：提取空间特征，通过Conv2d实现，参数包括输入通道数、输出通道数和卷积核大小。
  • 激活函数：使用ReLU引入非线性。
  • 池化层：通过max_pool2d降低特征图尺寸，减少计算量。
  • 全连接层：将特征映射到类别空间。

3.2 损失函数与优化器

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

选择依据：

  • 交叉熵损失：适用于多分类问题，直接比较预测概率分布与真实标签。
  • SGD优化器：加入动量项可加速收敛并减少震荡。

4. 模型训练与评估

4.1 训练循环

  1. def train(net, trainloader, criterion, optimizer, epochs=10):
  2.     for epoch in range(epochs):
  3.         running_loss = 0.0
  4.         for i, data in enumerate(trainloader, 0):
  5.             # 获取输入和标签
  6.             inputs, labels = data
  8.             # 梯度清零
  9.             optimizer.zero_grad()
  11.             # 前向传播
  12.             outputs = net(inputs)
  14.             # 计算损失
  15.             loss = criterion(outputs, labels)
  17.             # 反向传播
  18.             loss.backward()
  20.             # 更新参数
  21.             optimizer.step()
  23.             # 统计损失
  24.             running_loss += loss.item()
  25.             if i % 2000 == 1999:  # 每2000个batch打印一次
  26.                 print(f'Epoch {epoch + 1}, Batch {i + 1}, Loss: {running_loss / 2000:.3f}')
  27.                 running_loss = 0.0
  28.     print('Finished Training')
  30. train(net, trainloader, criterion, optimizer, epochs=10)

关键步骤：

  • optimizer.zero_grad()：清除历史梯度，防止梯度累积。
  • loss.backward()：自动计算梯度。
  • optimizer.step()：根据梯度更新参数。

4.2 模型评估

  1. def evaluate(net, testloader):
  2.     correct = 0
  3.     total = 0
  4.     with torch.no_grad():  # 禁用梯度计算
  5.         for data in testloader:
  6.             images, labels = data
  7.             outputs = net(images)
  8.             _, predicted = torch.max(outputs.data, 1)  # 获取概率最大的类别
  9.             total += labels.size(0)
  10.             correct += (predicted == labels).sum().item()
  12.     accuracy = 100 * correct / total
  13.     print(f'Accuracy on test set: {accuracy:.2f}%')
  14.     return accuracy
  16. evaluate(net, testloader)

评估指标：

  • 准确率：正确预测的样本数占总样本数的比例。

5. 完整代码与注释

  1. # 完整代码整合（含详细注释）
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5. import torch.optim as optim
  6. import torchvision
  7. import torchvision.transforms as transforms
  8. import matplotlib.pyplot as plt
  9. import numpy as np
  11. # 1. 数据准备
  12. transform = transforms.Compose([
  13.     transforms.ToTensor(),
  14.     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  15. ])
  17. trainset = torchvision.datasets.CIFAR10(
  18.     root='./data', train=True, download=True, transform=transform
  19. )
  20. trainloader = torch.utils.data.DataLoader(
  21.     trainset, batch_size=32, shuffle=True, num_workers=2
  22. )
  24. testset = torchvision.datasets.CIFAR10(
  25.     root='./data', train=False, download=True, transform=transform
  26. )
  27. testloader = torch.utils.data.DataLoader(
  28.     testset, batch_size=32, shuffle=False, num_workers=2
  29. )
  31. # 2. 定义网络
  32. class Net(nn.Module):
  33.     def __init__(self):
  34.         super(Net, self).__init__()
  35.         self.conv1 = nn.Conv2d(3, 6, 5)  # 输入3通道，输出6通道，5x5卷积核
  36.         self.conv2 = nn.Conv2d(6, 16, 5)
  37.         self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 全连接层
  38.         self.fc2 = nn.Linear(120, 84)
  39.         self.fc3 = nn.Linear(84, 10)
  41.     def forward(self, x):
  42.         x = F.max_pool2d(F.relu(self.conv1(x)), 2)  # 卷积+ReLU+池化
  43.         x = F.max_pool2d(F.relu(self.conv2(x)), 2)
  44.         x = x.view(-1, 16 * 5 * 5)  # 展平
  45.         x = F.relu(self.fc1(x))
  46.         x = F.relu(self.fc2(x))
  47.         x = self.fc3(x)
  48.         return x
  50. net = Net()
  52. # 3. 定义损失函数和优化器
  53. criterion = nn.CrossEntropyLoss()
  54. optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
  56. # 4. 训练函数
  57. def train(net, trainloader, criterion, optimizer, epochs=10):
  58.     for epoch in range(epochs):
  59.         running_loss = 0.0
  60.         for i, data in enumerate(trainloader, 0):
  61.             inputs, labels = data
  62.             optimizer.zero_grad()
  63.             outputs = net(inputs)
  64.             loss = criterion(outputs, labels)
  65.             loss.backward()
  66.             optimizer.step()
  67.             running_loss += loss.item()
  68.             if i % 2000 == 1999:
  69.                 print(f'Epoch {epoch + 1}, Batch {i + 1}, Loss: {running_loss / 2000:.3f}')
  70.                 running_loss = 0.0
  71.     print('Finished Training')
  73. train(net, trainloader, criterion, optimizer, epochs=10)
  75. # 5. 评估函数
  76. def evaluate(net, testloader):
  77.     correct = 0
  78.     total = 0
  79.     with torch.no_grad():
  80.         for data in testloader:
  81.             images, labels = data
  82.             outputs = net(images)
  83.             _, predicted = torch.max(outputs.data, 1)
  84.             total += labels.size(0)
  85.             correct += (predicted == labels).sum().item()
  86.     accuracy = 100 * correct / total
  87.     print(f'Accuracy on test set: {accuracy:.2f}%')
  88.     return accuracy
  90. evaluate(net, testloader)

6. 实践建议与扩展方向

  1. 模型改进：

    • 增加卷积层深度（如使用ResNet架构）。
    • 引入批归一化（nn.BatchNorm2d）加速收敛。
    • 使用数据增强（如随机裁剪、水平翻转）提升泛化能力。

  2. 超参数调优：

    • 学习率调度器（如torch.optim.lr_scheduler.StepLR）。
    • 网格搜索或贝叶斯优化寻找最优超参数。

  3. 部署应用：

    • 导出模型为TorchScript格式（torch.jit.trace）。
    • 部署至移动端（通过PyTorch Mobile）或云端服务。

总结

本文通过一个完整的案例，展示了使用PyTorch实现图像分类的全流程，包括数据加载、模型构建、训练优化和结果评估。代码中提供了详细的注释，帮助读者理解每个步骤的原理和实现细节。对于初学者，建议从简化模型开始，逐步尝试更复杂的架构和技巧；对于进阶用户，可探索分布式训练、模型压缩等高级主题。”

相关文章推荐

发表评论

最热文章

关于作者

  • 被阅读数
  • 被赞数
  • 被收藏数