pytorch cifar10分类

原创

mob64ca12f7ae31 2024-11-29 05:13:06 ©著作权

文章标签 数据集 python 深度学习 文章分类 PyTorch 人工智能

©著作权归作者所有：来自51CTO博客作者mob64ca12f7ae31的原创作品，请联系作者获取转载授权，否则将追究法律责任

PyTorch CIFAR-10 分类入门教程

计算机视觉是深度学习领域的重要应用之一，其中图像分类是最基础的任务之一。CIFAR-10数据集是图像分类领域中的经典数据集之一，由10类图像构成，是一个简单而有效的入门数据集。本文将介绍如何使用PyTorch库对CIFAR-10数据集进行分类，并提供相关代码示例。

CIFAR-10数据集简介

CIFAR-10数据集包含60,000张32x32像素的彩色图像，分为10类，每类有6,000张图像。具体类别如下：

飞机
汽车
鸟
猫
鹿
狗
青蛙
马
船
卡车

每条图像的标签范围在0到9之间，分别对应不同类别。

PyTorch简介

PyTorch是一个流行的深度学习框架，因其动态计算图和易于上手的API而备受欢迎。它非常适合研究和开发复杂的神经网络模型。

环境准备

在开始之前，你需要确保你的环境中安装了PyTorch和torchvision库。可以通过如下命令安装：

pip install torch torchvision

数据加载

我们将利用torchvision库中的CIFAR10类来下载和加载数据。首先，我们需要导入相关的库。

import torch
import torchvision
import torchvision.transforms as transforms

接下来，定义数据的转换方式，包括图像的归一化处理：

transform = transforms.Compose(
    [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=4,
                                          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=4,
                                         shuffle=False, num_workers=2)

定义神经网络

接下来，我们将定义一个简单的卷积神经网络（CNN）。CNN在图像处理任务中表现良好。下面是一个简单CNN模型的定义：

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, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

网络结构图

以下是这个CNN模型的结构图，我们可以用mermaid语法表示：

erDiagram
    CNN {
        string conv1
        string pool
        string conv2
        string fc1
        string fc2
        string fc3
    }
    CNN ||--o| Input : processes
    Input {
        string image
    }

训练神经网络

现在，我们可以开始训练模型。首先，我们需要定义损失函数和优化器。我们将使用交叉熵损失和SGD优化算法。

import torch.optim as optim

net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(2):  # 训练模型2个epoch
    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 % 2000 == 1999:  # 每2000个batch打印一次
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Finished Training')

测试模型

训练完成后，我们需要对模型进行测试。以下代码将计算在测试集上的准确率：

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'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')