PyTorch 系列教程:使用CNN实现图像分类

发布于:2025-03-13 ⋅ 阅读:(8) ⋅ 点赞:(0)

图像分类是计算机视觉领域的一项基本任务,也是深度学习技术的一个常见应用。近年来,卷积神经网络(cnn)和PyTorch库的结合由于其易用性和鲁棒性已经成为执行图像分类的流行选择。

理解卷积神经网络(cnn)

卷积神经网络是一类深度神经网络,对分析视觉图像特别有效。他们利用多层构建一个可以直接从图像中识别模式的模型。这些模型对于图像识别和分类等任务特别有用,因为它们不需要手动提取特征。

cnn的关键组成部分

  • 卷积层:这些层对输入应用卷积操作,将结果传递给下一层。每个过滤器(或核)可以捕获不同的特征,如边缘、角或其他模式。
  • 池化层:这些层减少了表示的空间大小,以减少参数的数量并加快计算速度。池化层简化了后续层的处理。
  • 完全连接层:在这些层中,神经元与前一层的所有激活具有完全连接,就像传统的神经网络一样。它们有助于对前一层识别的对象进行分类。
    在这里插入图片描述

使用PyTorch进行图像分类

PyTorch是开源的深度学习库,提供了极大的灵活性和多功能性。研究人员和从业人员广泛使用它来轻松有效地实现尖端的机器学习模型。

设置PyTorch

首先,确保在开发环境中安装了PyTorch。你可以通过pip安装它:

pip install torch torchvision

用PyTorch创建简单的CNN示例

下面是如何定义简单的CNN来使用PyTorch对图像进行分类的示例。

import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义CNN模型(修复了变量引用问题)
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)      # 第一个卷积层:3输入通道,6输出通道,5x5卷积核
        self.pool = nn.MaxPool2d(2, 2)        # 最大池化层:2x2窗口,步长2
        self.conv2 = nn.Conv2d(6, 16, 5)     # 第二个卷积层:6输入通道,16输出通道,5x5卷积核
        self.fc1 = nn.Linear(16 * 5 * 5, 120)# 全连接层1:400输入 -> 120输出
        self.fc2 = nn.Linear(120, 84)      # 全连接层2:120输入 -> 84输出
        self.fc3 = nn.Linear(84, 10)       # 输出层:84输入 -> 10类 logits

    def forward(self, x):
        # 输入形状:[batch_size, 3, 32, 32]
        x = self.pool(F.relu(self.conv1(x)))  # -> [batch, 6, 14, 14](池化后尺寸减半)
        x = self.pool(F.relu(self.conv2(x)))  # -> [batch, 16, 5, 5] 
        x = x.view(-1, 16 * 5 * 5)            # 展平为一维向量:16 * 5 * 5=400
        x = F.relu(self.fc1(x))             # -> [batch, 120]
        x = F.relu(self.fc2(x))             # -> [batch, 84]
        x = self.fc3(x)                     # -> [batch, 10](未应用softmax,配合CrossEntropyLoss使用)
        return x

这个特殊的网络接受一个输入图像,通过两组卷积和池化层,然后是三个完全连接的层。根据数据集的复杂性和大小调整网络的架构和超参数。

模型定义

  • SimpleCNN 继承自 nn.Module
  • 使用两个卷积层提取特征,三个全连接层进行分类
  • 最终输出未应用 softmax,而是直接输出 logits(与 CrossEntropyLoss 配合使用)

训练网络

对于训练,你需要一个数据集。PyTorch通过torchvision包提供了用于数据加载和预处理的实用程序。

import torchvision.transforms as transforms
import torchvision
from torch.utils.data import DataLoader

# 初始化模型、损失函数和优化器
net = SimpleCNN()               # 实例化模型
criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失函数(自动处理softmax)
optimizer = torch.optim.SGD(net.parameters(), 
                            lr=0.001,      # 学习率
                            momentum=0.9)   # 动量参数

# 数据预处理和加载
transform = transforms.Compose([
    transforms.ToTensor(),          
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  

# 加载CIFAR-10训练集
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True,
    download=True,  # 自动下载数据集
    transform=transform
)

trainloader = DataLoader(trainset, 
                     batch_size=4,   # 每个batch包含4张图像
                     shuffle=True)  # 打乱数据顺序

模型配置

  • 损失函数CrossEntropyLoss(自动包含 softmax 和 log_softmax)
  • 优化器:SGD with momentum,学习率 0.001

数据加载

  • 使用 torchvision.datasets.CIFAR10 加载数据集

  • batch_size:4(根据 GPU 内存调整,CIFAR-10 建议 batch size ≥ 32)

  • transforms.Compose 定义数据预处理流程:

    • ToTensor():将图像转换为 PyTorch Tensor
    • Normalize():标准化图像像素值到 [-1, 1]

加载数据后,训练过程包括通过数据集进行多次迭代,使用反向传播和合适的损失函数:

# 训练循环
for epoch in range(2):  # 进行2个epoch的训练
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        
        # 前向传播
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        optimizer.zero_grad()   # 清空梯度
        loss.backward()         # 计算梯度
        optimizer.step()       # 更新参数
        
        running_loss += loss.item()
        
        # 每2000个batch打印一次
        if i % 2000 == 1999:
            avg_loss = running_loss / 2000
            print(f'Epoch [{epoch+1}/{2}], Batch [{i+1}/2000], Loss: {avg_loss:.3f}')
            running_loss = 0.0

print("训练完成!")

训练循环

  • epoch:完整遍历数据集一次
  • batch:数据加载器中的一个批次
  • 梯度清零:每次反向传播前需要清空梯度
  • 损失计算outputs 的形状为 [batch_size, 10]labels 为整数标签

完整代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision
from torch.utils.data import DataLoader

# 定义CNN模型(修复了变量引用问题)
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)      # 第一个卷积层:3输入通道,6输出通道,5x5卷积核
        self.pool = nn.MaxPool2d(2, 2)        # 最大池化层:2x2窗口,步长2
        self.conv2 = nn.Conv2d(6, 16, 5)     # 第二个卷积层:6输入通道,16输出通道,5x5卷积核
        self.fc1 = nn.Linear(16 * 5 * 5, 120)# 全连接层1:400输入 -> 120输出
        self.fc2 = nn.Linear(120, 84)      # 全连接层2:120输入 -> 84输出
        self.fc3 = nn.Linear(84, 10)       # 输出层:84输入 -> 10类 logits

    def forward(self, x):
        # 输入形状:[batch_size, 3, 32, 32]
        x = self.pool(F.relu(self.conv1(x)))  # -> [batch, 6, 14, 14](池化后尺寸减半)
        x = self.pool(F.relu(self.conv2(x)))  # -> [batch, 16, 5, 5] 
        x = x.view(-1, 16 * 5 * 5)            # 展平为一维向量:16 * 5 * 5=400
        x = F.relu(self.fc1(x))             # -> [batch, 120]
        x = F.relu(self.fc2(x))             # -> [batch, 84]
        x = self.fc3(x)                     # -> [batch, 10](未应用softmax,配合CrossEntropyLoss使用)
        return x

# 初始化模型、损失函数和优化器
net = SimpleCNN()               # 实例化模型
criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失函数(自动处理softmax)
optimizer = torch.optim.SGD(net.parameters(), 
                            lr=0.001,      # 学习率
                            momentum=0.9)   # 动量参数

# 数据预处理和加载
transform = transforms.Compose([
    transforms.ToTensor(),            
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  
])

# 加载CIFAR-10训练集
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True,
    download=True,  # 自动下载数据集
    transform=transform
)
trainloader = DataLoader(trainset, 
                         batch_size=4,   # 每个batch包含4张图像
                         shuffle=True)  # 打乱数据顺序

# 训练循环
for epoch in range(2):  # 进行2个epoch的训练
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        
        # 前向传播
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        optimizer.zero_grad()   # 清空梯度
        loss.backward()         # 计算梯度
        optimizer.step()       # 更新参数
        
        running_loss += loss.item()
        
        # 每2000个batch打印一次
        if i % 2000 == 1999:
            avg_loss = running_loss / 2000
            print(f'Epoch [{epoch+1}/{2}], Batch [{i+1}/2000], Loss: {avg_loss:.3f}')
            running_loss = 0.0

print("训练完成!")

最后总结

通过PyTorch和卷积神经网络,你可以有效地处理图像分类任务。借助PyTorch的灵活性,可以根据特定的数据集和应用程序构建、训练和微调模型。示例代码仅为理论过程,实际项目中还有大量优化空间。


网站公告

今日签到

点亮在社区的每一天
去签到