Pytorch 使用CNN图像分类的实现

当涉及到图像分类时，卷积神经网络（CNN）是最常用的深度学习模型之一。在本攻略中，我们将介绍如何使用PyTorch实现CNN图像分类。我们将使用CIFAR-10数据集作为示例数据集。

步骤1：加载数据集

首先，我们需要加载CIFAR-10数据集。CIFAR-10数据集包含10个类别的60000个32x32彩色图像。我们将使用torchvision库中的CIFAR-10数据集来加载数据集。以下是一个示例代码，演示了如何加载CIFAR-10数据集：

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)

# 定义类别标签
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

在上面的代码中，我们首先定义了一个数据转换，该转换将图像转换为张量，并对图像进行归一化。我们使用torchvision.datasets.CIFAR10()函数加载训练集和测试集，并使用torch.utils.data.DataLoader()函数创建数据加载器。我们还定义了类别标签，该标签对应于CIFAR-10数据集中的10个类别。

步骤2：定义CNN模型

接下来，我们需要定义CNN模型。我们将使用PyTorch中的nn.Module类来定义CNN模型。以下是一个示例代码，演示了如何定义CNN模型：

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

# 定义CNN模型
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模型
net = Net()

在上面的代码中，我们首先定义了一个Net类，该类继承自nn.Module类。我们在Net类的构造函数中定义了CNN模型的各个层。我们使用nn.Conv2d()函数定义卷积层，使用nn.MaxPool2d()函数定义池化层，使用nn.Linear()函数定义全连接层。在Net类的forward()函数中，我们定义了CNN模型的前向传播过程。我们首先将输入x传递给第一个卷积层，然后使用ReLU激活函数和池化层。接下来，我们将输出传递给第二个卷积层，并再次使用ReLU激活函数和池化层。然后，我们将输出展平，并传递给三个全连接层。最后，我们返回输出。

步骤3：定义损失函数和优化器

接下来，我们需要定义损失函数和优化器。我们将使用交叉熵损失函数和随机梯度下降（SGD）优化器。以下是一个示例代码，演示了如何定义损失函数和优化器：

import torch.optim as optim

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

在上面的代码中，我们使用nn.CrossEntropyLoss()函数定义交叉熵损失函数。我们使用optim.SGD()函数定义SGD优化器，并将CNN模型的参数和学习率作为参数传递给该函数。

步骤4：训练CNN模型

现在，我们已经准备好训练CNN模型了。我们将使用训练集来训练CNN模型，并使用测试集来评估CNN模型的性能。以下是一个示例代码，演示了如何训练CNN模型：

# 训练CNN模型
for epoch in range(2):  # 多次循环遍历数据集
    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个小批量数据打印一次统计信息
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

在上面的代码中，我们使用两个for循环来遍历训练集。在内部循环中，我们首先获取输入和标签。然后，我们使用optimizer.zero_grad()函数清除梯度。接下来，我们将输入传递给CNN模型，并计算输出和损失。然后，我们使用loss.backward()函数计算梯度，并使用optimizer.step()函数更新CNN模型的参数。最后，我们打印统计信息，包括损失和当前的迭代次数。

步骤5：测试CNN模型

最后，我们需要使用测试集来测试CNN模型的性能。以下是一个示例代码，演示了如何测试CNN模型：

# 测试CNN模型
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('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))

在上面的代码中，我们使用torch.no_grad()上下文管理器来禁用梯度计算。在循环中，我们首先获取输入和标签。然后，我们将输入传递给CNN模型，并计算输出。接下来，我们使用torch.max()函数获取预测结果，并使用(predicted == labels).sum().item()函数统计正确预测的数量和总数量。最后，我们打印测试结果，包括CNN模型的准确率。

示例1：使用GPU训练CNN模型

如果您的计算机具有GPU，您可以使用GPU来加速CNN模型的训练。以下是一个示例代码，演示了如何使用GPU训练CNN模型：

# 将CNN模型移动到GPU上
net.cuda()

# 将输入和标签移动到GPU上
inputs, labels = inputs.cuda(), labels.cuda()

# 计算输出和损失
outputs = net(inputs)
loss = criterion(outputs, labels)

# 计算梯度并更新参数
loss.backward()
optimizer.step()

在上面的代码中，我们首先使用net.cuda()函数将CNN模型移动到GPU上。然后，我们使用inputs.cuda()和labels.cuda()函数将输入和标签移动到GPU上。接下来，我们将输入传递给CNN模型，并计算输出和损失。然后，我们使用loss.backward()函数计算梯度，并使用optimizer.step()函数更新CNN模型的参数。

示例2：使用预训练模型进行迁移学习

如果您的数据集与CIFAR-10数据集不同，您可以使用预训练模型进行迁移学习。以下是一个示例代码，演示了如何使用预训练模型进行迁移学习：

import torchvision.models as models

# 加载预训练模型
resnet18 = models.resnet18(pretrained=True)

# 将预训练模型的最后一层替换为新的全连接层
num_ftrs = resnet18.fc.in_features
resnet18.fc = nn.Linear(num_ftrs, 10)

# 实例化CNN模型
net = resnet18

在上面的代码中，我们首先使用models.resnet18()函数加载预训练模型。然后，我们使用nn.Linear()函数将预训练模型的最后一层替换为新的全连接层。最后，我们实例化CNN模型，并将预训练模型作为其基础。这样，我们就可以使用预训练模型的特征提取能力来训练新的CNN模型。

本站文章如无特殊说明，均为本站原创，如若转载，请注明出处：Pytorch 使用CNN图像分类的实现 - Python技术站