在本教程中,我们将使用PyTorch框架来实现一个简单的卷积神经网络(CNN),用于识别MNIST数据集中的手写数字。MNIST数据集是一个广泛使用的数据集,包含60,000个训练样本和10,000个测试样本,每个样本是一个28x28的灰度图像,表示0到9之间的手写数字。

1. 导入必要的库

首先,我们需要导入一些必要的库,包括PyTorch、torchvision、matplotlib等。torchvision提供了对MNIST数据集的便捷访问,而matplotlib用于可视化数据。

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

2. 设置设备

我们将使用GPU来加速训练过程(如果可用)。通过以下代码,我们可以检查是否有可用的GPU,并将其设置为默认设备。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

3. 数据预处理与加载

接下来,我们需要对数据进行预处理,并将其加载到DataLoader中。我们将使用transforms.Compose来定义一系列的数据转换操作,包括将图像转换为张量并进行归一化。

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_set = datasets.MNIST(root='D:\\python_text\\python\\pytorch深度学习实战\\dataset', transform=transform, train=True, download=True)
test_set = datasets.MNIST(root='D:\\python_text\\python\\pytorch深度学习实战\\dataset', transform=transform, train=False, download=True)

batch_size = 64
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=4)

4. 可视化数据

在训练模型之前,我们可以先可视化一些训练数据,以确保数据加载正确。

plt.figure(figsize=(20,10))
for i in range(20):
    plt.subplot(5,10,i+1)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)
    plt.imshow(train_set.data[i], cmap=plt.cm.binary)
    plt.xlabel(train_set.targets[i])

plt.show()

5. 定义模型

我们将定义一个简单的卷积神经网络模型。该模型包含两个卷积层、两个最大池化层和两个全连接层。

from torch import nn, optim

class Model(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(64*7*7, 64)
        self.fc2 = nn.Linear(64, 10)
        self.relu = nn.ReLU()
        self.flatten = nn.Flatten()

    def forward(self, x):
        x = self.pool1(self.relu(self.conv1(x)))
        x = self.pool2(self.relu(self.conv2(x)))
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        return self.fc2(x)

model = Model()
model.to(device)

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

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

criterion = nn.CrossEntropyLoss()
criterion.to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)

7. 训练模型

我们定义了一个训练函数,用于训练模型并在每个epoch结束后测试模型的准确率。

def train(model, train_loader, test_loader, optimizer, criterion, num_epochs, device):
    for epoch in range(num_epochs):
        # 训练模型
        model.train()
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            output = model(x)
            loss = criterion(output, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print('epoch:', epoch+1, 'loss:', loss.item())

        # 测试模型
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for x, y in test_loader:
                x, y = x.to(device), y.to(device)
                output = model(x)
                _, predicted = torch.max(output.data, dim=1)
                correct += (predicted == y).sum().item()
                total += y.size(0)
            print('acc:', correct/total)

train(model, train_loader, test_loader, optimizer, criterion, 10, device)

8. 结果

在训练过程中,我们可以看到每个epoch的损失值和测试集上的准确率。经过10个epoch的训练,模型的准确率可以达到98%以上。

9. 总结

通过本教程,我们学习了如何使用PyTorch构建一个简单的卷积神经网络来识别MNIST手写数字。我们涵盖了数据加载、模型定义、训练和测试等步骤。希望本教程对你理解如何使用PyTorch进行深度学习有所帮助。

如果你有任何问题或建议,欢迎在评论区留言!

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐