零基础AI入门：PyTorch手把手教你搭建神经网络模型识别手写数字

2024-01-05 11:29:05

在人工智能领域，深度学习模型的应用越来越广泛，尤其是在图像识别方面。手写数字识别作为深度学习的一个经典应用，不仅能够锻炼我们的动手能力，还能加深对神经网络的理解。本文将详细介绍如何使用PyTorch搭建一个神经网络模型，用于识别手写数字。

PyTorch简介

PyTorch是一个开源的机器学习库，特别适用于深度学习的应用。它的设计理念是动态计算图，这使得模型构建和调试更加直观和灵活。PyTorch支持Python编程语言，提供了丰富的API，方便开发者快速上手。

PyTorch的安装

要开始使用PyTorch，首先需要确保你的系统已经安装了Python和pip。然后，你可以通过以下命令安装PyTorch：

pip install torch torchvision

安装完成后，可以通过以下代码验证PyTorch是否安装成功：

import torch
print(torch.__version__)

搭建神经网络

在深度学习中，神经网络是核心组件。一个简单的神经网络通常包括输入层、隐藏层和输出层。以手写数字识别为例，我们可以使用一个包含一个隐藏层的两层神经网络。

首先，导入必要的库：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

接下来，定义一个简单的神经网络结构：

class SimpleNeuralNetwork(nn.Module):
    def __init__(self):
        super(SimpleNeuralNetwork, self).__init__()
        self.fc1 = nn.Linear(784, 128)  # 输入层到隐藏层
        self.fc2 = nn.Linear(128, 10)   # 隐藏层到输出层

    def forward(self, x):
        x = x.view(-1, 784)          # 将图像展平为一维数组
        x = torch.relu(self.fc1(x))  # 使用ReLU激活函数
        x = self.fc2(x)              # 输出层没有激活函数
        return x

在这个网络中，nn.Linear用于定义全连接层，relu是一个常用的激活函数。

训练神经网络

训练神经网络需要准备数据集和标签，并使用优化器进行参数更新。以下是完整的训练代码：

# 定义模型
model = SimpleNeuralNetwork()

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

# 加载MNIST数据集
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 创建DataLoader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

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