PyTorch--卷积神经网络（CNN）模型实现手写数字识别

前言

今天要介绍的这段代码是一个使用PyTorch框架实现的卷积神经网络（CNN）模型，用于对MNIST数据集进行分类的示例。MNIST数据集是手写数字识别领域的一个标准数据集，包含0到9的灰度图像。

代码的主要组成部分如下：

1. 导入必要的库：导入PyTorch、PyTorch神经网络模块、torchvision（用于处理图像数据集）和transforms（用于图像预处理）。

2. 设备配置：设置模型运行的设备，优先使用GPU（如果可用），否则使用CPU。

3. 超参数设置：定义了训练迭代的轮数(num_epochs)、类别数(num_classes)、批次大小(batch_size)和学习率(learning_rate)。

4. 加载MNIST数据集：使用torchvision.datasets.MNIST加载MNIST训练集和测试集，并应用transforms.ToTensor将图像转换为张量。

5. 创建数据加载器：使用torch.utils.data.DataLoader创建训练和测试数据的加载器，以便在训练和测试过程中批量加载数据。

6. 定义卷积神经网络模型：定义了一个名为ConvNet的类，继承自nn.Module。模型包含两个卷积层（每层后接批量归一化和ReLU激活函数），以及一个全连接层。

7. 实例化模型并移动到设备：创建ConvNet模型的实例，并将其移动到之前设置的设备上。

8. 定义损失函数和优化器：使用nn.CrossEntropyLoss作为损失函数，torch.optim.Adam作为优化器。

9. 训练模型：进行多个epoch的训练，每个epoch中对数据集进行遍历，执行前向传播、损失计算、反向传播和参数更新。

10. 测试模型：在测试阶段，将模型设置为评估模式，并禁用梯度计算以提高效率，然后计算模型在测试集上的准确率。

11. 保存模型：使用torch.save保存训练后的模型参数到文件，以便将来可以重新加载和使用模型。

完整代码

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


# Device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Hyper parameters
num_epochs = 5
num_classes = 10
batch_size = 100
learning_rate = 0.001

# MNIST dataset
train_dataset = torchvision.datasets.MNIST(root='../../data/',
                                           train=True, 
                                           transform=transforms.ToTensor(),
                                           download=True)

test_dataset = torchvision.datasets.MNIST(root='../../data/',
                                          train=False, 
                                          transform=transforms.ToTensor())

# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size, 
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size, 
                                          shuffle=False)

# Convolutional neural network (two convolutional layers)
class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7*7*32, num_classes)
        
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

model = ConvNet(num_classes).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Train the model
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        
        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

# Test the model
model.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item(