Python 实现卷积神经网络图像分类

卷积神经网络（Convolutional Neural Network，CNN）是一种用于处理图像数据的深度学习模型。它通过卷积层、池化层和全连接层的组合，可以自动学习图像中的特征，从而实现图像分类。后续介绍使用 Python 和 PyTorch 库实现卷积神经网络图像分类。

Solomon_肖哥弹架构

476人浏览 · 2025-02-17 00:56:39

Solomon_肖哥弹架构 · 2025-02-17 00:56:39 发布

在这里插入图片描述
卷积神经网络（Convolutional Neural Network，CNN）是一种用于处理图像数据的深度学习模型。它通过卷积层、池化层和全连接层的组合，可以自动学习图像中的特征，从而实现图像分类。后续介绍使用 Python 和 PyTorch 库实现卷积神经网络图像分类。

一、背景介绍

CIFAR-10 是一个广泛使用的图像分类数据集，包含 10 个类别的彩色图像，每个类别有 6000 张图像。这些图像的大小为 32x32 像素，每个像素有三个颜色通道（RGB）。数据集分为 50,000 张训练图像和 10,000 张测试图像。

二、数据加载

数据集可以通过 PyTorch 提供的 torchvision.datasets 模块加载。我们将使用 CIFAR10 数据集，并对数据进行预处理，包括归一化和张量转换。

import torch
import torchvision
import torchvision.transforms as transforms

# 定义数据预处理
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.RandomCrop(32, padding=4),  # 随机裁剪
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化
])

# 加载训练数据
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
                                          shuffle=True, num_workers=2)

# 加载测试数据
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64,
                                         shuffle=False, num_workers=2)

三、模型构建

我们构建一个简单的卷积神经网络，包含两个卷积层、两个池化层和两个全连接层。

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 卷积层
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)  # 输入通道数为3，输出通道数为32，卷积核大小为3x3
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)  # 输入通道数为32，输出通道数为64
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)  # 输入通道数为64，输出通道数为128
        
        # 池化层
        self.pool = nn.MaxPool2d(2, 2)  # 最大池化层，池化窗口为2x2，步幅为2
        
        # 全连接层
        self.fc1 = nn.Linear(128 * 4 * 4, 512)  # 输入节点数为128*4*4，输出节点数为512
        self.fc2 = nn.Linear(512, 10)  # 输入节点数为512，输出节点数为10

    def forward(self, x):
        # 前向传播
        x = F.relu(self.conv1(x))  # 第一层卷积 + ReLU激活
        x = self.pool(x)  # 最大池化

        x = F.relu(self.conv2(x))  # 第二层卷积 + ReLU激活
        x = self.pool(x)  # 最大池化

        x = F.relu(self.conv3(x))  # 第三层卷积 + ReLU激活
        x = self.pool(x)  # 最大池化

        x = x.view(-1, 128 * 4 * 4)  # 将张量展平

        x = F.relu(self.fc1(x))  # 第一层全连接 + ReLU激活
        x = self.fc2(x)  # 第二层全连接
        return x

# 创建模型实例
net = Net()

四、模型训练

我们使用交叉熵损失函数和 Adam 优化器来训练模型。

import torch.optim as optim

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

# 训练模型
for epoch in range(10):  # 训练10个周期
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()  # 清空梯度
        outputs = net(inputs)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新权重

        running_loss += loss.item()
        if i % 200 == 199:  # 每200个批次打印一次
            print(f'[Epoch {epoch + 1}, Batch {i + 1}] loss: {running_loss / 200:.3f}')
            running_loss = 0.0

print('训练完成')

五、模型评估

训练完成后，我们在测试集上评估模型的性能。

correct = 0
total = 0
with torch.no_grad():  # 在评估过程中关闭梯度计算
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'测试集准确率: {100 * correct / total:.2f}%')