PyTorch使用CNN实现图像分类

更新时间：2025年03月13日 08:49:47 作者：梦想画家

图像分类是计算机视觉领域的一项基本任务,也是深度学习技术的一个常见应用,近年来,卷积神经网络（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

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

模型定义：

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

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

您可能感兴趣的文章:

使用Python中PDB模块中的命令来调试Python代码的教程
这篇文章主要介绍了使用Python中PDB模块中的命令来调试Python代码的教程,包括设置断点来修改代码等、对于Python团队项目工作有一定帮助，需要的朋友可以参考下
2015-03-03
Python使用DrissionPage实现网页自动化采集
DrissionPage 是一个基于 python 的网页自动化工具,它既能控制浏览器,也能收发数据包,还能把两者合而为一,可兼顾浏览器自动化的便利性和 requests 的高效率,本文给大家介绍了Python使用DrissionPage实现网页自动化采集,需要的朋友可以参考下
2025-03-03
Django 导出项目依赖库到 requirements.txt过程解析
这篇文章主要介绍了Django 导出项目依赖库到 requirements.txt过程解析,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下
2019-08-08
将tensorflow模型打包成PB文件及PB文件读取方式
今天小编就为大家分享一篇将tensorflow模型打包成PB文件及PB文件读取方式，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2020-01-01
Gradio机器学习模型快速部署工具quickstart前篇
这篇文章主要为大家介绍了Gradio机器学习模型快速部署工具quickstart准备原文翻译，有需要的朋友可以借鉴参考下，希望能够有所帮助，祝大家多多进步，早日升职加薪
2023-04-04
Python利用CNN实现对时序数据进行分类
这篇文章主要为大家详细介绍了Python如何利用CNN实现对时序数据进行分类功能，文中的示例代码讲解详细，感兴趣的小伙伴可以了解一下
2023-02-02
python开发微信服务号消息推送示例
这篇文章主要为大家介绍了python开发微信服务号消息推送示例详解,有需要的朋友可以借鉴参考下,希望能够有所帮助,祝大家多多进步,早日升职加薪
2023-10-10
Django数据库迁移的实现步骤
本文主要介绍了Django数据库迁移的实现步骤,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧
2025-08-08
Python实现计算文件MD5和SHA1的方法示例
这篇文章主要介绍了Python实现计算文件MD5和SHA1的方法,结合具体实例形式分析了Python针对文件MD5及SHA1的计算方法,需要的朋友可以参考下
2019-06-06
python的open函数使用案例代码
python打开文件使用open()函数，返回一个指向文件的指针，在python中使用open函数对文件进行处理，这篇文章主要介绍了python中的open函数使用,需要的朋友可以参考下
2023-02-02