10 从零到上手系统学习 PyTorch - 构建简单的神经网络

在这一小节中，我们将详细介绍如何使用 PyTorch 构建一个简单的神经网络。我们将分步骤进行，包括数据准备、模型定义、训练及评估等部分。

1. 数据准备

在构建神经网络之前，首先需要准备训练和测试数据。我们将使用 torchvision 库中的 MNIST 数据集，这是一个手写数字识别的数据集。

1.1 导入必要的库

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

1.2 下载和加载数据集

我们将使用 torchvision.datasets 中的 MNIST 数据集，并应用一些基本的转换，比如将图像转换为张量并进行归一化处理。

# 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize((0.5,), (0.5,))  # 归一化处理
])

# 下载训练集和测试集
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

2. 定义神经网络

接下来，我们将定义一个简单的全连接神经网络。我们将创建一个子类 Net 继承自 nn.Module。

2.1 创建神经网络模型

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)  # 输入层到隐藏层
        self.fc2 = nn.Linear(128, 64)        # 隐藏层到隐藏层
        self.fc3 = nn.Linear(64, 10)         # 隐藏层到输出层
        self.relu = nn.ReLU()                 # 激活函数

    def forward(self, x):
        x = x.view(-1, 28 * 28)  # 将28x28的图像展平成784维的向量
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x

3. 定义损失函数和优化器

选择一个适当的损失函数和优化器对于训练模型是至关重要的。对于多类别分类问题，我们通常使用 CrossEntropyLoss 作为损失函数，使用 Adam 作为优化器。

# 实例化模型、损失函数和优化器
model = Net()
criterion = nn.CrossEntropyLoss()  # 损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 优化器

4. 训练模型

现在，我们可以开始训练模型。我们将迭代多个 epoch，每个 epoch 中使用训练数据进行模型的更新。

4.1 训练过程

num_epochs = 5  # 训练轮数

for epoch in range(num_epochs):
    for inputs, labels in trainloader:
        optimizer.zero_grad()  # 清空之前的梯度
        outputs = model(inputs)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新权重

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

5. 评估模型

训练完成后，我们需要评估模型在测试集上的效果。我们将计算模型的准确率。

5.1 测试过程

correct = 0
total = 0

with torch.no_grad():  # 不需要计算梯度
    for inputs, labels in testloader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)  # 获取预测结果
        total += labels.size(0)  # 更新总样本数
        correct += (predicted == labels).sum().item()  # 统计预测正确的样本数

print(f'Accuracy: {100 * correct / total:.2f}%')

总结

通过上述步骤，我们成功实现了一个简单的神经网络，并对其进行训练和评估。在这个过程中，我们使用了 PyTorch 的基本功能，包括数据加载、模型构建、损失计算和梯度更新等。接下来，您可以尝试调整网络结构、优化器或者其他超参数，进一步提高模型的性能。