python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > PyTorch CNN图像分类

PyTorch使用CNN实现图像分类

作者：梦想画家

图像分类是计算机视觉领域的一项基本任务,也是深度学习技术的一个常见应用,近年来,卷积神经网络（cnn）和PyTorch库的结合由于其易用性和鲁棒性已经成为执行图像分类的流行选择,所以本文给大家介绍了PyTorch使用CNN实现图像分类的示例,需要的朋友可以参考下

理解卷积神经网络（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

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

模型定义：

训练网络

对于训练，你需要一个数据集。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)  # 打乱数据顺序

模型配置：

数据加载：

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

# 训练循环
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("训练完成！")

训练循环：

完整代码

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的灵活性，可以根据特定的数据集和应用程序构建、训练和微调模型。示例代码仅为理论过程，实际项目中还有大量优化空间。

以上就是PyTorch使用CNN实现图像分类的详细内容，更多关于PyTorch CNN图像分类的资料请关注脚本之家其它相关文章！

您可能感兴趣的文章:
阅读全文