PyTorch 手写数字识别：构建简单高效的神经网络模型

import torch
import torch.nn as nn
import os
from torch.autograd import Variable
from torchvision.datasets import MNIST
from torchvision.transforms import transforms
from torch.utils.data import DataLoader
from torch.optim import Adam
import numpy

# 设置 GPU ID
os.environ['CUDA_VISIBLE_DEVICES'] = '1'


# 构建神经网络
class Unit(nn.Module):
    def __init__(self, in_features, out_features):
        super(Unit, self).__init__()

        self.fc = nn.Linear(in_features, out_features, True)
        self.relu = nn.ReLU()

    def forward(self, input):
        output = self.fc(input)
        output = self.relu(output)
        return output


class SimpleNet(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleNet, self).__init__()

        self.unit1 = Unit(in_features=784, out_features=1024)
        self.unit2 = Unit(in_features=1024, out_features=1024)
        self.unit3 = Unit(in_features=1024, out_features=2048)
        self.net = nn.Sequential(self.unit1, self.unit2, self.unit3)
        self.fc = nn.Linear(in_features=2048, out_features=num_classes)

    def forward(self, input):
        output = self.net(input)
        output = self.fc(output)
        return output


# 创建学习率调整函数，每 30 个 epoch 将学习率除以 10
def adjust_learning_rate(optimizer, epoch):
    lr = 0.001
    if epoch > 180:
        lr = lr / 1000000
    elif epoch > 150:
        lr = lr / 100000
    elif epoch > 120:
        lr = lr / 10000
    elif epoch > 90:
        lr = lr / 1000
    elif epoch > 60:
        lr = lr / 100
    elif epoch > 30:
        lr = lr / 10
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


# 训练过程
def train(model, train_acc, train_loss, train_loader, optimizer, loss_fn, cuda_avail=True):
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        if cuda_avail:
            images = Variable(images.cuda())
            images = images.reshape(-1, 784)
            labels = Variable(labels.cuda())
        optimizer.zero_grad()
        outputs = model(images)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.cpu().data.numpy() * images.size(0)
        _, prediction = torch.max(outputs.data, 1)
        train_acc += torch.sum(prediction == labels.data)
    train_acc = train_acc / 60000
    train_loss = train_loss / 60000
    return train_acc, train_loss


# 在训练时测试模型
def test(model, test_loader, cuda_avail=True):
    model.eval()
    test_acc = 0.0
    for i, (images, labels) in enumerate(test_loader):
        if cuda_avail:
            images = Variable(images.cuda())
            images = images.reshape(-1, 784)
            labels = Variable(labels.cuda())

        # 使用测试集中的图像预测类别
        outputs = model(images)
        _, prediction = torch.max(outputs.data, 1)
        test_acc += torch.sum(prediction == labels.data)

    # 计算所有 10000 张测试图像的平均准确率和损失
    test_acc = test_acc / 10000
    return test_acc


# 保存权重
def save_models(model, epoch):
    torch.save(model.state_dict(), 'MNIST_model_{}.model'.format(epoch))
    print('Chekcpoint saved')


def main():
    # 1. 加载训练数据集和测试数据集
    train_transformations = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    train_set = MNIST(root='./data', train=True, transform=train_transformations, download=True)
    train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=4)
    test_transformations = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    test_set = MNIST(root='./data', train=False, transform=test_transformations, download=True)
    test_loader = DataLoader(test_set, batch_size=128, shuffle=False, num_workers=0)
    # 2. 初始化模型
    cuda_avail = torch.cuda.is_available()
    model = SimpleNet(num_classes=10)
    cuda_avail = torch.cuda.is_available()
    if cuda_avail:
        model.cuda()
    # 3. 定义优化器和损失函数
    optimizer = Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
    loss_fn = nn.CrossEntropyLoss()
    # 4. 训练
    num_epochs = 100
    best_acc = 0.0
    for epoch in range(num_epochs):
        train_acc = 0.0
        train_loss = 0.0
        train_acc, train_loss = train(model, train_acc, train_loss, train_loader,
                                      optimizer, loss_fn, cuda_avail)
        adjust_learning_rate(optimizer, epoch)
        test_acc = test(model, test_loader, cuda_avail)
        if test_acc > best_acc:
            save_models(model, epoch)
            best_acc = test_acc
            print('Epoch {}, Train Accuracy: {} , TrainLoss: {} , Test Accuracy: {}'.format(
                epoch, train_acc, train_loss, test_acc))


if __name__ == '__main__':
    main()

这段代码使用 PyTorch 构建了一个简单的神经网络模型，用于识别手写数字。

代码解释：

导入必要的库:
- torch: 用于构建和训练神经网络的核心库。
- torch.nn: 包含用于构建神经网络的模块和类。
- os: 用于设置环境变量，例如指定使用的 GPU。
- torch.autograd: 提供自动微分功能，用于计算梯度。
- torchvision: 提供用于图像处理的数据集、模型和预训练权重。
- torch.utils.data: 提供数据加载和预处理的工具。
- torch.optim: 包含各种优化算法，例如 Adam 优化器。
- numpy: 用于数值计算的库。
定义神经网络模型:
- Unit: 一个简单的线性层，后跟 ReLU 激活函数。
- SimpleNet: 由多个 Unit 层组成的序列模型，最后连接一个全连接层，用于分类。
定义学习率调整函数:
- adjust_learning_rate: 根据训练的 epoch 数量动态调整学习率，以提高训练效果。
定义训练和测试函数:
- train: 训练模型，计算损失并更新模型参数。
- test: 在测试集上评估模型性能。
定义保存模型函数:
- save_models: 保存训练过程中性能最佳的模型权重。
主函数 (main):
- 加载 MNIST 数据集，并进行预处理。
- 初始化模型、优化器和损失函数。
- 进行多轮训练和测试，并根据测试准确率保存最佳模型。

代码总结:

这段代码展示了如何使用 PyTorch 构建一个简单的手写数字识别模型。它涵盖了数据加载、模型定义、训练和测试等关键步骤，并提供了一个完整的示例，可以作为学习 PyTorch 的起点。