PyTorch CNN实战之MNIST手写数字识别示例
本文将介绍如何使用PyTorch实现一个卷积神经网络（CNN）来对MNIST手写数字数据集进行分类，即识别手写数字。
MNIST数据集是一个经典的图像分类数据集，其中包含有手写数字图片和标签，图片尺寸为28 x 28像素。该数据集共有10个类别：数字0-9。
模型架构
我们将使用三层卷积和池化层，以及两层全连接层来构建我们的CNN模型。
首先，我们定义一个名为MNISTModel的类，继承自torch.nn.Module。在类中，我们定义了模型的各个层，并在forward()函数中实现前向传播过程。代码如下：
import torch
import torch.nn as nn
class MNISTModel(nn.Module):
def __init__(self):
super(MNISTModel, self).__init__()
self.conv1 = nn.Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
self.relu1 = nn.ReLU()
self.maxpool1 = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
self.conv2 = nn.Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
self.relu2 = nn.ReLU()
self.maxpool2 = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
self.conv3 = nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
self.relu3 = nn.ReLU()
self.maxpool3 = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
self.fc1 = nn.Linear(64 * 3 * 3, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.conv1(x)
x = self.relu1(x)
x = self.maxpool1(x)
x = self.conv2(x)
x = self.relu2(x)
x = self.maxpool2(x)
x = self.conv3(x)
x = self.relu3(x)
x = self.maxpool3(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
return x
接下来，我们定义训练函数和测试函数。
训练函数
我们定义了一个train()函数来训练我们的模型。首先，我们定义了优化器和损失函数。在本例中，我们使用交叉熵损失函数（CrossEntropyLoss）和Adam优化器。
然后，我们进行多次训练迭代。在每次迭代中，我们首先将数据和标签放入设备（GPU或CPU）中。接着，我们将输入数据通过模型进行前向传播，计算损失函数。然后，我们使用反向传播和优化器对模型进行优化。
最后，我们计算本次迭代的训练损失和准确率，并更新 测试准确率和相应的 模型参数。代码如下：
def train(model, device, train_loader, optimizer, epoch):
model.train()
train_loss = 0.0
train_correct_count = 0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = nn.CrossEntropyLoss()(output, target)
loss.backward()
optimizer.step()
train_loss += loss.item() * data.size(0)
pred = output.argmax(dim=1, keepdim=True)
train_correct_count += pred.eq(target.view_as(pred)).sum().item()
train_loss /= len(train_loader.dataset)
train_accuracy = 100.0 * train_correct_count / len(train_loader.dataset)
print('Train Epoch: {} \tLoss: {:.6f} \tAccuracy: {:.2f}%'.format(
epoch, train_loss, train_accuracy))
return train_accuracy
测试函数
我们定义了一个test()函数来测试我们的模型在测试集上的表现。我们在函数中计算模型在所有测试数据上的准确率，最终返回测试准确率和相应的测试损失。代码如下：
def test(model, device, test_loader):
model.eval()
test_loss = 0.0
test_correct_count = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
loss = nn.CrossEntropyLoss()(output, target)
test_loss += loss.item() * data.size(0)
pred = output.argmax(dim=1, keepdim=True)
test_correct_count += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
test_accuracy = 100.0 * test_correct_count / len(test_loader.dataset)
print('Test Loss: {:.6f} \tAccuracy: {:.2f}%'.format(
test_loss, test_accuracy))
return test_accuracy, test_loss
模型训练
接下来，我们将使用上述定义的MNISTModel、train()函数和test()函数来训练和测试我们的模型。
首先，我们需要下载MNIST数据集，并将其转换为torch.Tensor格式。我们可以使用torchvision.datasets库中的MNIST数据集加载器来完成这个任务。
import torchvision.datasets as datasets from torchvision.transforms import ToTensor train_set = datasets.MNIST(root='./data', train=True, download=True, transform=ToTensor()) test_set = datasets.MNIST(root='./data', train=False, download=False, transform=ToTensor())
接下来，我们将加载器传入torch.utils.data.DataLoader来生成可迭代的数据集。我们可以指定batch_size、shuffle和num_workers参数来控制批量大小、是否打乱数据集和使用的CPU内核数。在本例中，我们使用批量大小为64，打乱数据集，并使用4个CPU内核来加快数据加载。
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True, num_workers=4) test_loader = torch.utils.data.DataLoader(test_set, batch_size=64, shuffle=False, num_workers=4)
接下来，我们定义了一些训练超参数（如学习率、训练次数等）和一些辅助变量（如 测试准确率、对应的 模型参数等）。我们还定义了一个device变量，将根据可用的GPU或CPU来自动选择设备。如果GPU可用，则使用GPU进行训练和测试。否则，我们将使用CPU。
learning_rate = 0.001
num_epochs = 10
best_test_accuracy = 0.0
best_model_params = None
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
最后，我们定义了一个模型实例，并将其移动到指定的设备上。然后，我们使用定义的优化器和学习率对模型进行优化。在每个epoch中，我们对训练集进行一轮训练，并在测试集上对模型进行一次测试。如果模型在测试集上表现更好，则更新相应的参数。
model = MNISTModel().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(1, num_epochs + 1):
train_accuracy = train(model, device, train_loader, optimizer, epoch)
test_accuracy, test_loss = test(model, device, test_loader)
if test_accuracy > best_test_accuracy:
best_test_accuracy = test_accuracy
best_model_params = model.state_dict()
print('Best Test Accuracy: {:.2f}%'.format(best_test_accuracy))
输出结果
训练完毕后，我们可以输出 测试准确率并使用 模型参数来预测新数据。
下面是本次实验的输出结果：
`
Train Epoch: 1 Loss: 0.394917 Accuracy: 87.29%
Test