PyTorch 入门：利用标准网络结构探索图像识别

人工智能

2024-01-04 15:03:56

前言

图像识别作为计算机视觉领域的重要分支，在诸多应用中扮演着至关重要的角色。而 PyTorch 作为深度学习框架中的佼佼者，凭借其灵活性、高效性和广泛的网络结构支持，为图像识别任务提供了得天独厚的优势。

本文将带领您踏上图像识别的征程，以 PyTorch 作为我们的工具，利用标准网络结构 ResNet18，深入探索构建和训练图像识别模型的过程。

PyTorch 简介

PyTorch 是一个 Python 第一的深度学习框架，以其动态计算图机制而著称。该机制允许我们在运行时灵活地构建和修改计算图，为研究人员和开发者提供了高度的灵活性。

ResNet18：标准网络结构

ResNet18 是一种卷积神经网络（CNN），由 He 等人于 2015 年提出。它属于 ResNet 家族，是一种使用残差连接的深度 CNN。与传统的 CNN 相比，ResNet 在深层网络中表现出更高的准确性和训练稳定性。

ResNet18 由 18 个卷积层、池化层和激活函数组成，能够有效地提取图像特征并进行分类。其深度和残差连接使其能够学习复杂的高级特征，从而提升图像识别性能。

数据集准备

我们将使用 CIFAR10 数据集进行模型训练。CIFAR10 包含 60,000 张 32x32 像素的彩色图像，分为 10 个类别（飞机、汽车、鸟类、猫、鹿、狗、青蛙、马、船、卡车）。

使用 PyTorch 内置的 torchvision.datasets 模块可以轻松加载和预处理 CIFAR10 数据集。

import torchvision
import torch.utils.data as data

transform = torchvision.transforms.ToTensor()

train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)

test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)

train_loader = data.DataLoader(train_dataset, batch_size=64, shuffle=True)

test_loader = data.DataLoader(test_dataset, batch_size=64, shuffle=False)

模型构建

利用 PyTorch 搭建 ResNet18 模型非常简单，我们可以直接从 torchvision.models 模块导入预先训练好的 ResNet18 模型。

import torchvision.models as models

model = models.resnet18(pretrained=False)

模型训练

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

import torch.nn as nn
import torch.optim as optim

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10

for epoch in range(num_epochs):
    for i, (inputs, labels) in enumerate(train_loader):

        outputs = model(inputs)

        loss = criterion(outputs, labels)

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        if i % 100 == 0:

            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item()}')

模型评估

训练完成后，我们可以使用测试集来评估模型的性能。

model.eval()

with torch.no_grad():

    correct = 0

    total = 0

    for images, labels in test_loader:

        outputs = model(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} %')