PyTorch入门教程—手写数字图像分类

PyTorch 入门教程

PyTorch 是一个开源的机器学习框架，可以用于实现深度神经网络模型。本教程将介绍如何使用 PyTorch 进行简单的神经网络构建和训练任务。

安装 PyTorch

在开始教程之前，需要先安装 PyTorch。可以通过以下命令安装 PyTorch：

pip install torch torchvision

构建简单神经网络模型

我们将构建一个简单的神经网络模型来分类手写数字图像。首先，导入必要的 PyTorch 模块和数据集。

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

# Define the transform to normalize the data
# 定义数据预处理方法
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,))])

# Download and load the MNIST dataset
# 下载和加载 MNIST 手写数字数据集
trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
testset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

接下来，我们定义一个简单的神经网络模型。该模型有两个全连接层和一个输出层。为了增加模型的非线性特征，我们在每个层之间使用 ReLU 激活函数。

# Define the neural network
# 定义神经网络模型
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) # 输出层

    def forward(self, x):
        x = x.view(-1, 28 * 28) # 将输入数据展开成一维向量
        x = nn.functional.relu(self.fc1(x)) # ReLU 激活函数，第一层网络
        x = nn.functional.relu(self.fc2(x)) # ReLU 激活函数，第二层网络
        x = self.fc3(x) # 输出层
        return x

net = Net()

训练神经网络模型

现在我们的模型已经定义好了，我们需要对它进行训练。我们使用交叉熵损失函数和随机梯度下降优化器来训练模型。

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

# Train the neural network
# 训练神经网络模型
for epoch in range(5):
    running_loss = 0.0
    for i, data in enumerate(trainset, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 1000 == 999:    # 每1000个数据样本批次打印一次运行损失
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 1000))
            running_loss = 0.0

print('Finished Training')

测试神经网络模型

训练完成后，我们需要对模型进行测试，看看它在新数据上的性能如何。我们可以在测试集上计算模型的准确率。

# Test the neural network
# 测试神经网络模型
correct = 0
total = 0
with torch.no_grad():
    for data in testset:
        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))

结论

在本教程中，我们使用 PyTorch 构建了一个简单的神经网络模型，并对其进行训练和测试。通过这个例子，您可以了解 PyTorch 的基本用法。希望本教程对您有所帮助。