用PyTorch轻松搭建AlexNet图像分类器，一探深度学习奥妙

2023-04-19 22:51:36

图像分类：使用 PyTorch 和 AlexNet 构建一个图像分类器

图像分类是计算机视觉领域中至关重要的任务之一，它能够识别和分类各种物体，在医疗、交通和安防等领域有着广泛的应用。在众多经典的深度学习模型中，AlexNet 因其出色的图像分类能力而备受推崇。本文将手把手地指导您使用 PyTorch 框架和 AlexNet 模型，从零开始构建一个图像分类器，带您踏入深度学习的精彩世界。

AlexNet：计算机视觉的里程碑

AlexNet 于 2012 年由 Alex Krizhevsky 等人提出，是计算机视觉领域的一座里程碑。它由 8 层神经网络组成，包括 5 个卷积层和 3 个全连接层。通过优化网络结构和权重，AlexNet 在图像分类任务上取得了突破性的表现。

PyTorch：构建图像分类器的理想平台

PyTorch 是一个功能强大的深度学习框架，以其灵活性、易用性和高性能著称。PyTorch 提供了丰富的工具和模块，让您可以轻松地构建和训练深度学习模型。我们将使用 PyTorch 来构建我们的图像分类器，并逐步介绍卷积层、池化层、全连接层等基本模块的使用方法。

构建 AlexNet 模型

import torch
import torch.nn as nn

class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()

        # 第一层卷积
        self.conv1 = nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2)

        # 第一层池化
        self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2)

        # 第二层卷积
        self.conv2 = nn.Conv2d(64, 192, kernel_size=5, padding=2)

        # 第二层池化
        self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2)

        # 第三层卷积
        self.conv3 = nn.Conv2d(192, 384, kernel_size=3, padding=1)

        # 第四层卷积
        self.conv4 = nn.Conv2d(384, 256, kernel_size=3, padding=1)

        # 第五层卷积
        self.conv5 = nn.Conv2d(256, 256, kernel_size=3, padding=1)

        # 第一层全连接层
        self.fc1 = nn.Linear(256 * 6 * 6, 4096)

        # 第二层全连接层
        self.fc2 = nn.Linear(4096, 4096)

        # 第三层全连接层
        self.fc3 = nn.Linear(4096, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.maxpool1(x)
        x = self.conv2(x)
        x = self.maxpool2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

准备数据

我们使用 CIFAR-10 数据集进行训练和测试。该数据集包含 6 万张彩色图像，分为 10 个类别（如汽车、船只、飞机等）。

import torchvision.datasets as datasets
import torchvision.transforms as transforms

train_data = datasets.CIFAR10(root='./data', train=True,
                            download=True, transform=transforms.ToTensor())

test_data = datasets.CIFAR10(root='./data', train=False,
                           download=True, transform=transforms.ToTensor())

定义损失函数和优化器

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

训练模型

for epoch in range(10):
    for i, (inputs, labels) in enumerate(train_loader):
        # 前向传播
        outputs = model(inputs)

        # 计算损失
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()

        # 更新权重
        optimizer.step()

        # 清空梯度
        optimizer.zero_grad()

评估模型

correct = 0
total = 0
with torch.no_grad():
    for i, (inputs, labels) in enumerate(test_loader):
        outputs = model(inputs)
        _, 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))