logo

开发者热搜

使用PyTorch构建图像分类系统：完整代码与深度解析

作者：渣渣辉2025.09.19 17:05浏览量：2

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

使用PyTorch构建图像分类系统：完整代码与深度解析

图像分类是计算机视觉领域的核心任务之一，PyTorch作为主流深度学习框架，提供了简洁高效的API支持。本文将通过完整代码示例，展示如何使用PyTorch实现从数据准备到模型部署的全流程，所有代码均包含详细注释，确保读者能够理解每个步骤的实现原理。

一、环境准备与依赖安装

首先需要安装PyTorch及相关依赖库。推荐使用conda创建虚拟环境：

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

关键依赖说明：

  • torch：PyTorch核心库
  • torchvision：提供计算机视觉常用数据集和模型架构
  • matplotlib：用于可视化训练过程
  • numpy：数值计算基础库

二、数据集准备与预处理

1. 使用CIFAR-10数据集

CIFAR-10包含10个类别的60000张32x32彩色图像，分为50000张训练集和10000张测试集。

  1. import torchvision
  2. import torchvision.transforms as transforms
  4. # 定义数据预处理流程
  5. transform = transforms.Compose([
  6.     transforms.ToTensor(),  # 将PIL图像转换为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. # 加载测试集
  25. testset = torchvision.datasets.CIFAR10(
  26.     root='./data', 
  27.     train=False,
  28.     download=True, 
  29.     transform=transform
  30. )
  31. testloader = torch.utils.data.DataLoader(
  32.     testset, 
  33.     batch_size=32,
  34.     shuffle=False, 
  35.     num_workers=2
  36. )
  38. # 类别名称
  39. classes = ('plane', 'car', 'bird', 'cat', 'deer', 
  40.            'dog', 'frog', 'horse', 'ship', 'truck')

关键点解析：

  • transforms.Compose：组合多个数据预处理操作
  • ToTensor()：将HWC格式的PIL图像转换为CHW格式的Tensor
  • Normalize：使用均值和标准差进行标准化，这里使用(0.5,0.5,0.5)将像素值映射到[-1,1]区间
  • DataLoader：实现批量加载、数据打乱和多线程加载

2. 自定义数据集加载

对于自定义数据集，可以继承torch.utils.data.Dataset类：

  1. from torch.utils.data import Dataset
  2. import os
  3. from PIL import Image
  5. class CustomImageDataset(Dataset):
  6.     def __init__(self, img_dir, transform=None):
  7.         self.img_labels = []
  8.         self.img_paths = []
  9.         self.transform = transform
  11.         # 遍历目录，假设子目录名为类别名
  12.         for class_name in os.listdir(img_dir):
  13.             class_path = os.path.join(img_dir, class_name)
  14.             if os.path.isdir(class_path):
  15.                 for img_name in os.listdir(class_path):
  16.                     self.img_paths.append(os.path.join(class_path, img_name))
  17.                     self.img_labels.append(classes.index(class_name))
  19.     def __len__(self):
  20.         return len(self.img_paths)
  22.     def __getitem__(self, idx):
  23.         img_path = self.img_paths[idx]
  24.         image = Image.open(img_path)
  25.         label = self.img_labels[idx]
  27.         if self.transform:
  28.             image = self.transform(image)
  30.         return image, label

三、模型架构设计

1. 基础CNN模型

  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.         # 输入通道3(RGB)，输出通道32，3x3卷积核
  8.         self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
  9.         self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
  10.         self.pool = nn.MaxPool2d(2, 2)  # 2x2最大池化
  11.         self.fc1 = nn.Linear(64 * 8 * 8, 512)  # CIFAR-10经过两次池化后为8x8
  12.         self.fc2 = nn.Linear(512, 10)  # 10个输出类别
  13.         self.dropout = nn.Dropout(0.25)
  15.     def forward(self, x):
  16.         # 第一层卷积+ReLU+池化
  17.         x = self.pool(F.relu(self.conv1(x)))
  18.         # 第二层卷积+ReLU+池化
  19.         x = self.pool(F.relu(self.conv2(x)))
  20.         # 展平特征图
  21.         x = x.view(-1, 64 * 8 * 8)
  22.         # 全连接层+ReLU+Dropout
  23.         x = self.dropout(F.relu(self.fc1(x)))
  24.         # 输出层
  25.         x = self.fc2(x)
  26.         return x

架构解析：

  • 两个卷积层提取空间特征，每个卷积层后接ReLU激活函数和最大池化
  • 两个全连接层完成分类，中间加入Dropout防止过拟合
  • 输入32x32x3图像，经过两次2x2池化后变为8x8x64特征图

2. 使用预训练模型

PyTorch提供了多种预训练模型，可通过torchvision.models加载：

  1. import torchvision.models as models
  3. def get_pretrained_model(model_name='resnet18', pretrained=True, num_classes=10):
  4.     if model_name == 'resnet18':
  5.         model = models.resnet18(pretrained=pretrained)
  6.         # 修改最后一层全连接网络
  7.         num_ftrs = model.fc.in_features
  8.         model.fc = nn.Linear(num_ftrs, num_classes)
  9.     elif model_name == 'vgg16':
  10.         model = models.vgg16(pretrained=pretrained)
  11.         num_ftrs = model.classifier[6].in_features
  12.         model.classifier[6] = nn.Linear(num_ftrs, num_classes)
  13.     else:
  14.         raise ValueError("Unsupported model name")
  15.     return model

四、训练流程实现

1. 完整训练代码

  1. import torch
  2. import torch.optim as optim
  3. from tqdm import tqdm  # 进度条库
  5. def train_model(model, trainloader, testloader, criterion, optimizer, num_epochs=10):
  6.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  7.     model.to(device)
  9.     for epoch in range(num_epochs):
  10.         # 训练阶段
  11.         model.train()
  12.         running_loss = 0.0
  13.         correct = 0
  14.         total = 0
  16.         # 使用tqdm显示进度条
  17.         train_loop = tqdm(trainloader, desc=f'Epoch {epoch+1}/{num_epochs}')
  18.         for inputs, labels in train_loop:
  19.             inputs, labels = inputs.to(device), labels.to(device)
  21.             # 梯度清零
  22.             optimizer.zero_grad()
  24.             # 前向传播
  25.             outputs = model(inputs)
  26.             loss = criterion(outputs, labels)
  28.             # 反向传播和优化
  29.             loss.backward()
  30.             optimizer.step()
  32.             # 统计信息
  33.             running_loss += loss.item()
  34.             _, predicted = torch.max(outputs.data, 1)
  35.             total += labels.size(0)
  36.             correct += (predicted == labels).sum().item()
  38.             # 更新进度条信息
  39.             train_loop.set_postfix(loss=running_loss/(train_loop.n+1), 
  40.                                   acc=100.*correct/total)
  42.         # 测试阶段
  43.         test_loss, test_acc = evaluate_model(model, testloader, criterion, device)
  44.         print(f'Epoch {epoch+1}, Train Loss: {running_loss/len(trainloader):.4f}, '
  45.               f'Train Acc: {100*correct/total:.2f}%, '
  46.               f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%')
  48. def evaluate_model(model, testloader, criterion, device):
  49.     model.eval()
  50.     test_loss = 0.0
  51.     correct = 0
  52.     total = 0
  54.     with torch.no_grad():
  55.         for inputs, labels in testloader:
  56.             inputs, labels = inputs.to(device), labels.to(device)
  57.             outputs = model(inputs)
  58.             loss = criterion(outputs, labels)
  60.             test_loss += loss.item()
  61.             _, predicted = torch.max(outputs.data, 1)
  62.             total += labels.size(0)
  63.             correct += (predicted == labels).sum().item()
  65.     return test_loss/len(testloader), 100*correct/total
  67. # 初始化模型
  68. model = SimpleCNN()
  69. # 或者使用预训练模型
  70. # model = get_pretrained_model('resnet18')
  72. # 定义损失函数和优化器
  73. criterion = nn.CrossEntropyLoss()
  74. optimizer = optim.Adam(model.parameters(), lr=0.001)
  76. # 开始训练
  77. train_model(model, trainloader, testloader, criterion, optimizer, num_epochs=10)

2. 关键训练参数说明

  • 学习率：控制参数更新步长，常用值为0.001(Adam)或0.01(SGD)
  • 批量大小：影响内存使用和梯度估计稳定性，CIFAR-10常用32或64
  • 优化器选择：
    • Adam：自适应学习率，收敛快
    • SGD+Momentum：可能获得更好泛化性能
  • 损失函数：分类任务通常使用交叉熵损失

五、模型评估与可视化

1. 混淆矩阵实现

  1. import numpy as np
  2. import matplotlib.pyplot as plt
  3. from sklearn.metrics import confusion_matrix
  4. import seaborn as sns
  6. def plot_confusion_matrix(model, testloader, classes, device):
  7.     model.eval()
  8.     all_labels = []
  9.     all_preds = []
  11.     with torch.no_grad():
  12.         for inputs, labels in testloader:
  13.             inputs, labels = inputs.to(device), labels.to(device)
  14.             outputs = model(inputs)
  15.             _, predicted = torch.max(outputs, 1)
  16.             all_labels.extend(labels.cpu().numpy())
  17.             all_preds.extend(predicted.cpu().numpy())
  19.     cm = confusion_matrix(all_labels, all_preds)
  20.     plt.figure(figsize=(10,8))
  21.     sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
  22.                 xticklabels=classes, yticklabels=classes)
  23.     plt.xlabel('Predicted')
  24.     plt.ylabel('True')
  25.     plt.title('Confusion Matrix')
  26.     plt.show()
  28. # 调用示例
  29. plot_confusion_matrix(model, testloader, classes, device)

2. 训练过程可视化

  1. def plot_training_curve(train_losses, test_losses, train_accs, test_accs):
  2.     plt.figure(figsize=(12, 4))
  4.     plt.subplot(1, 2, 1)
  5.     plt.plot(train_losses, label='Train Loss')
  6.     plt.plot(test_losses, label='Test Loss')
  7.     plt.xlabel('Epoch')
  8.     plt.ylabel('Loss')
  9.     plt.legend()
  11.     plt.subplot(1, 2, 2)
  12.     plt.plot(train_accs, label='Train Accuracy')
  13.     plt.plot(test_accs, label='Test Accuracy')
  14.     plt.xlabel('Epoch')
  15.     plt.ylabel('Accuracy (%)')
  16.     plt.legend()
  18.     plt.tight_layout()
  19.     plt.show()
  21. # 需要在训练过程中记录这些指标
  22. # 示例数据
  23. epochs = range(1, 11)
  24. train_losses = [2.3, 1.8, 1.5, 1.2, 1.0, 0.9, 0.8, 0.7, 0.6, 0.5]
  25. test_losses = [2.1, 1.7, 1.4, 1.1, 0.95, 0.85, 0.78, 0.72, 0.68, 0.65]
  26. train_accs = [45, 58, 65, 70, 75, 78, 80, 82, 84, 85]
  27. test_accs = [50, 62, 68, 72, 76, 78, 80, 81, 82, 83]
  29. plot_training_curve(train_losses, test_losses, train_accs, test_accs)

六、模型部署建议

  1. 模型导出：使用torch.save保存模型参数

    1. torch.save(model.state_dict(), 'cifar_classifier.pth')

  2. 推理脚本示例：

    1. def predict_image(image_path, model, transform, classes, device):
    2.     image = Image.open(image_path)
    3.     image = transform(image).unsqueeze(0).to(device)
    4.     model.eval()
    5.     with torch.no_grad():
    6.         output = model(image)
    7.         _, predicted = torch.max(output.data, 1)
    8.     return classes[predicted.item()]

  3. 性能优化技巧：

    • 使用混合精度训练(torch.cuda.amp)
    • 模型量化减少内存占用
    • 使用TensorRT加速推理

七、常见问题解决方案

  1. 训练不收敛：

    • 检查学习率是否过大
    • 确认数据预处理是否正确
    • 尝试不同的优化器

  2. 过拟合问题：

    • 增加数据增强
    • 添加Dropout层
    • 使用L2正则化

  3. GPU内存不足：

    • 减小批量大小
    • 使用梯度累积
    • 清理缓存(torch.cuda.empty_cache())

本文完整代码可在GitHub获取，建议读者从简单CNN开始实践，逐步尝试预训练模型和更复杂的架构。通过调整超参数和观察训练曲线，可以深入理解深度学习模型的工作原理。

相关文章推荐

发表评论

最热文章

关于作者

  • 被阅读数
  • 被赞数
  • 被收藏数
活动
咨询