logo

开发者热搜

使用PyTorch实现图像分类：完整代码与深度解析

作者：热心市民鹿先生2025.09.18 17:43浏览量：0

简介：本文详细介绍如何使用PyTorch框架实现图像分类任务，包含从数据加载到模型训练的全流程代码，并附有详细注释。适合PyTorch初学者和有一定基础的开发者参考。

使用PyTorch实现图像分类：完整代码与深度解析

引言

图像分类是计算机视觉领域的核心任务之一，广泛应用于人脸识别、医学影像分析、自动驾驶等场景。PyTorch作为主流深度学习框架，以其动态计算图和简洁的API设计，成为实现图像分类任务的理想选择。本文将通过一个完整的CIFAR-10分类案例，详细展示如何使用PyTorch实现图像分类，包含所有关键代码和详细注释。

环境准备

在开始之前，请确保已安装以下环境：

  • Python 3.8+
  • PyTorch 1.12+
  • torchvision 0.13+
  • NumPy 1.21+
  • Matplotlib 3.5+

推荐使用conda创建虚拟环境：

  1. conda create -n pytorch_cls python=3.8
  2. conda activate pytorch_cls
  3. pip install torch torchvision numpy matplotlib

1. 数据准备与预处理

1.1 数据集加载

我们使用torchvision内置的CIFAR-10数据集，该数据集包含10个类别的60000张32x32彩色图像。

  1. import torch
  2. from torchvision import datasets, transforms
  3. from torch.utils.data import DataLoader
  5. # 定义数据预处理流程
  6. transform = transforms.Compose([
  7.     transforms.ToTensor(),  # 将PIL图像或numpy数组转为Tensor，并缩放到[0,1]
  8.     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化到[-1,1]
  9. ])
  11. # 加载训练集和测试集
  12. train_dataset = datasets.CIFAR10(
  13.     root='./data', 
  14.     train=True,
  15.     download=True,
  16.     transform=transform
  17. )
  18. test_dataset = datasets.CIFAR10(
  19.     root='./data',
  20.     train=False,
  21.     download=True,
  22.     transform=transform
  23. )
  25. # 创建数据加载器
  26. batch_size = 64
  27. train_loader = DataLoader(
  28.     train_dataset,
  29.     batch_size=batch_size,
  30.     shuffle=True,  # 每个epoch打乱数据
  31.     num_workers=2  # 使用2个子进程加载数据
  32. )
  33. test_loader = DataLoader(
  34.     test_dataset,
  35.     batch_size=batch_size,
  36.     shuffle=False,
  37.     num_workers=2
  38. )

关键点说明：

  • transforms.Compose：将多个预处理操作组合在一起
  • ToTensor()：自动将HWC格式的图像转为CHW格式的Tensor
  • Normalize：使用(mean, std)参数进行标准化，这里使用CIFAR-10的均值和标准差
  • DataLoader：提供批量加载、多线程加载和打乱数据等功能

1.2 数据可视化

为了验证数据加载是否正确，我们可以可视化部分训练样本：

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

2. 模型定义

2.1 基础CNN模型

我们定义一个简单的CNN模型，包含3个卷积层和2个全连接层：

  1. import torch.nn as nn
  2. import torch.nn.functional as F
  4. class SimpleCNN(nn.Module):
  5.     def __init__(self):
  6.         super(SimpleCNN, 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.         # CIFAR-10图像原始大小32x32，经过3次池化(每次尺寸减半)后为4x4
  16.         # 128个通道，所以展平后维度为128*4*4=2048
  17.         self.fc1 = nn.Linear(128 * 4 * 4, 512)
  18.         self.fc2 = nn.Linear(512, 10)  # 10个类别
  20.         # 最大池化层
  21.         self.pool = nn.MaxPool2d(2, 2)
  23.     def forward(self, x):
  24.         # 卷积1 -> ReLU -> 池化
  25.         x = self.pool(F.relu(self.conv1(x)))
  26.         # 卷积2 -> ReLU -> 池化
  27.         x = self.pool(F.relu(self.conv2(x)))
  28.         # 卷积3 -> ReLU -> 池化
  29.         x = self.pool(F.relu(self.conv3(x)))
  31.         # 展平特征图
  32.         x = x.view(-1, 128 * 4 * 4)
  34.         # 全连接层
  35.         x = F.relu(self.fc1(x))
  36.         x = self.fc2(x)
  37.         return x

模型结构说明：

  1. 3个卷积层，每层后接ReLU激活和2x2最大池化
  2. 卷积核大小均为3x3，使用padding=1保持空间尺寸
  3. 两个全连接层，第一个有512个神经元，第二个输出10个类别
  4. 输入图像32x32，经过3次池化后变为4x4

2.2 更先进的模型（可选）

对于更高性能的需求，可以使用ResNet等现代架构：

  1. from torchvision import models
  3. def get_resnet18(pretrained=False):
  4.     model = models.resnet18(pretrained=pretrained)
  5.     # 修改第一个卷积层以适应3通道输入（ResNet默认用于ImageNet的3通道）
  6.     # 实际上resnet18已经支持3通道，这里只是示例
  7.     # 修改最后的全连接层以适应10个类别
  8.     num_ftrs = model.fc.in_features
  9.     model.fc = nn.Linear(num_ftrs, 10)
  10.     return model

3. 训练过程

3.1 初始化模型和优化器

  1. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  2. print(f"Using device: {device}")
  4. model = SimpleCNN().to(device)
  5. criterion = nn.CrossEntropyLoss()  # 交叉熵损失
  6. optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # Adam优化器

3.2 训练函数

  1. def train_model(model, criterion, optimizer, num_epochs=10):
  2.     for epoch in range(num_epochs):
  3.         model.train()  # 设置为训练模式
  4.         running_loss = 0.0
  5.         correct = 0
  6.         total = 0
  8.         for i, (inputs, labels) in enumerate(train_loader):
  9.             inputs, labels = inputs.to(device), labels.to(device)
  11.             # 梯度清零
  12.             optimizer.zero_grad()
  14.             # 前向传播
  15.             outputs = model(inputs)
  16.             loss = criterion(outputs, labels)
  18.             # 反向传播和优化
  19.             loss.backward()
  20.             optimizer.step()
  22.             # 统计信息
  23.             running_loss += loss.item()
  24.             _, predicted = torch.max(outputs.data, 1)
  25.             total += labels.size(0)
  26.             correct += (predicted == labels).sum().item()
  28.             # 每100个batch打印一次统计信息
  29.             if i % 100 == 99:
  30.                 print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], '
  31.                       f'Loss: {running_loss/100:.3f}, Acc: {100*correct/total:.2f}%')
  32.                 running_loss = 0.0
  34.         # 每个epoch结束后计算并打印测试集准确率
  35.         test_acc = test_model(model, test_loader)
  36.         print(f'Epoch [{epoch+1}/{num_epochs}] completed, Test Acc: {test_acc:.2f}%')
  38. def test_model(model, test_loader):
  39.     model.eval()  # 设置为评估模式
  40.     correct = 0
  41.     total = 0
  42.     with torch.no_grad():  # 不计算梯度
  43.         for inputs, labels in test_loader:
  44.             inputs, labels = inputs.to(device), labels.to(device)
  45.             outputs = model(inputs)
  46.             _, predicted = torch.max(outputs.data, 1)
  47.             total += labels.size(0)
  48.             correct += (predicted == labels).sum().item()
  49.     return 100 * correct / total

3.3 启动训练

  1. num_epochs = 10
  2. train_model(model, criterion, optimizer, num_epochs)

4. 模型评估与保存

4.1 评估模型

训练完成后，我们可以在测试集上评估模型性能：

  1. test_acc = test_model(model, test_loader)
  2. print(f'Final Test Accuracy: {test_acc:.2f}%')

4.2 保存模型

  1. # 保存整个模型（包括结构和参数）
  2. torch.save(model.state_dict(), 'cifar10_cnn.pth')
  4. # 加载模型的代码示例
  5. def load_model():
  6.     model = SimpleCNN().to(device)
  7.     model.load_state_dict(torch.load('cifar10_cnn.pth'))
  8.     model.eval()
  9.     return model

5. 完整代码整合

以下是整合后的完整代码：

  1. # 完整实现PyTorch图像分类
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5. import torch.optim as optim
  6. from torchvision import datasets, transforms, models
  7. from torch.utils.data import DataLoader
  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. train_dataset = datasets.CIFAR10(
  18.     root='./data', train=True, download=True, transform=transform
  19. )
  20. test_dataset = datasets.CIFAR10(
  21.     root='./data', train=False, download=True, transform=transform
  22. )
  24. batch_size = 64
  25. train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
  26. test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)
  28. # 2. 模型定义
  29. class SimpleCNN(nn.Module):
  30.     def __init__(self):
  31.         super(SimpleCNN, self).__init__()
  32.         self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
  33.         self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
  34.         self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
  35.         self.fc1 = nn.Linear(128 * 4 * 4, 512)
  36.         self.fc2 = nn.Linear(512, 10)
  37.         self.pool = nn.MaxPool2d(2, 2)
  39.     def forward(self, x):
  40.         x = self.pool(F.relu(self.conv1(x)))
  41.         x = self.pool(F.relu(self.conv2(x)))
  42.         x = self.pool(F.relu(self.conv3(x)))
  43.         x = x.view(-1, 128 * 4 * 4)
  44.         x = F.relu(self.fc1(x))
  45.         x = self.fc2(x)
  46.         return x
  48. # 3. 训练设置
  49. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  50. model = SimpleCNN().to(device)
  51. criterion = nn.CrossEntropyLoss()
  52. optimizer = optim.Adam(model.parameters(), lr=0.001)
  54. # 4. 训练和测试函数
  55. def train_model(model, criterion, optimizer, num_epochs=10):
  56.     for epoch in range(num_epochs):
  57.         model.train()
  58.         running_loss = 0.0
  59.         correct = 0
  60.         total = 0
  62.         for i, (inputs, labels) in enumerate(train_loader):
  63.             inputs, labels = inputs.to(device), labels.to(device)
  65.             optimizer.zero_grad()
  66.             outputs = model(inputs)
  67.             loss = criterion(outputs, labels)
  68.             loss.backward()
  69.             optimizer.step()
  71.             running_loss += loss.item()
  72.             _, predicted = torch.max(outputs.data, 1)
  73.             total += labels.size(0)
  74.             correct += (predicted == labels).sum().item()
  76.             if i % 100 == 99:
  77.                 print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], '
  78.                       f'Loss: {running_loss/100:.3f}, Acc: {100*correct/total:.2f}%')
  79.                 running_loss = 0.0
  81.         test_acc = test_model(model, test_loader)
  82.         print(f'Epoch [{epoch+1}/{num_epochs}] completed, Test Acc: {test_acc:.2f}%')
  84. def test_model(model, test_loader):
  85.     model.eval()
  86.     correct = 0
  87.     total = 0
  88.     with torch.no_grad():
  89.         for inputs, labels in test_loader:
  90.             inputs, labels = inputs.to(device), labels.to(device)
  91.             outputs = model(inputs)
  92.             _, predicted = torch.max(outputs.data, 1)
  93.             total += labels.size(0)
  94.             correct += (predicted == labels).sum().item()
  95.     return 100 * correct / total
  97. # 5. 启动训练
  98. num_epochs = 10
  99. train_model(model, criterion, optimizer, num_epochs)
  101. # 6. 保存模型
  102. torch.save(model.state_dict(), 'cifar10_cnn.pth')

6. 性能优化建议

  1. 数据增强：在训练时使用随机裁剪、水平翻转等增强技术

    1. transform_train = transforms.Compose([
    2.     transforms.RandomHorizontalFlip(),
    3.     transforms.RandomCrop(32, padding=4),
    4.     transforms.ToTensor(),
    5.     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    6. ])

  2. 学习率调度：使用学习率衰减策略

    1. scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
    2. # 在每个epoch后调用scheduler.step()

  3. 批归一化：在卷积层后添加批归一化层

    1. self.bn1 = nn.BatchNorm2d(32)
    2. # 在forward中使用
    3. x = self.pool(F.relu(self.bn1(self.conv1(x))))

  4. 使用更先进的模型：如ResNet、EfficientNet等预训练模型

7. 总结与展望

本文详细介绍了使用PyTorch实现图像分类的完整流程，包括数据准备、模型定义、训练过程和模型评估。通过CIFAR-10数据集的实践，读者可以掌握：

  • PyTorch的基本数据加载机制
  • CNN模型的设计原则
  • 训练循环的实现细节
  • 模型评估和保存的方法

未来工作可以探索：

  1. 更大规模的数据集（如ImageNet）
  2. 更复杂的模型架构（如Transformer）
  3. 分布式训练以加速模型收敛
  4. 模型压缩技术以部署到移动端

希望本文能为PyTorch初学者提供实用的参考，帮助快速上手图像分类任务。

相关文章推荐

发表评论

最热文章

关于作者

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