Python 使用 PyTorch 的 nn 模块实现 LeNet 手写数字识别

以下是使用 PyTorch 的 nn 模块实现 LeNet 模型的代码，其中包括一些注释：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

# 定义 LeNet 模型
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # 卷积层 1，输入为 1 通道，输出为 6 通道，卷积核大小为 5*5
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        # 池化层 1，池化核大小为 2*2
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 卷积层 2，输入为 6 通道，输出为 16 通道，卷积核大小为 5*5
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        # 池化层 2，池化核大小为 2*2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 全连接层 1，输入为 16*5*5，输出为 120
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        # 全连接层 2，输入为 120，输出为 84
        self.fc2 = nn.Linear(120, 84)
        # 输出层，输入为 84，输出为 10
        self.fc3 = nn.Linear(84, 10)

    # 前向传播
    def forward(self, x):
        # 卷积-激活-池化-卷积-激活-池化-全连接-激活-全连接-激活-输出
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 使用 LeNet 模型进行手写数字识别
def train_lenet():
    # 加载 MNIST 数据集
    train_data = torch.utils.data.DataLoader(
        torch.utils.data.datasets.MNIST('mnist_data', train=True, download=True,
                       transform=torch.utils.data.transforms.Compose([
                           torch.utils.data.transforms.ToTensor(),
                           torch.utils.data.transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=64, shuffle=True)

    # 定义 LeNet 模型
    net = LeNet()

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

    # 训练模型
    for epoch in range(10):
        running_loss = 0.0
        for i, data in enumerate(train_data, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 100 == 99:
                print('[%d, %5d] loss: %.3f' % 
                      (epoch + 1, i + 1, running_loss / 100))
                running_loss = 0.0

    print('Finished Training')
    

if __name__ == '__main__':
    train_lenet()

在这个例子中，我们使用了 PyTorch 的 nn 模块来实现 LeNet 模型，并在 MNIST 数据集上进行训练和测试。在训练过程中，我们使用了交叉熵损失函数和随机梯度下降优化器。