导入所需库

import numpy as np import torch

导入 pytorch 内置的 mnist 数据

from torchvision.datasets import mnist

导入预处理模块

import torchvision.transforms as transforms from torch.utils.data import DataLoader

导入nn及优化器

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

定义超参数

train_batch_size = 16 test_batch_size = 16 learning_rate = 0.01 num_epoches = 20

定义预处理函数

transform = transforms.Compose([transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.RandomVerticalFlip(), # 随机垂直翻转 transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])

下载数据，并对数据进行预处理

train_dataset = mnist.MNIST('../data/', train=True, # 是否为训练集 transform=transform, # 数据预处理 download=False) # 是否下载 test_dataset = mnist.MNIST('../data/', train=False, transform=transform)

得到一个生成器

train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True) # 数据随机打乱 test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False)

构建神经网络模型

class Net(nn.Module): def init(self, in_dim, n_hidden_1, n_hidden_2, out_dim): super(Net, self).init() # 展开数据 self.flatten = nn.Flatten() # 第一层全连接和Batch Normalization self.layer1 = nn.Sequential(nn.Linear(in_dim, n_hidden_1), nn.BatchNorm1d(n_hidden_1)) # 第二层全连接和Batch Normalization self.layer2 = nn.Sequential(nn.Linear(n_hidden_1, n_hidden_2), nn.BatchNorm1d(n_hidden_2)) # 输出层 self.out = nn.Sequential(nn.Linear(n_hidden_2, out_dim)) # Dropout层 self.dropout = nn.Dropout(0.5)

def forward(self, x):
    x = self.flatten(x)     # 展开数据
    x = F.relu(self.layer1(x))      # 使用ReLU激活函数进行非线性转换
    x = self.dropout(x)     # 使用Dropout层进行正则化
    x = F.relu(self.layer2(x))
    x = self.dropout(x)
    x = F.softmax(self.out(x), dim=1)       # 使用Softmax进行激活
    return x

实例化模型，并将其移动到GPU上进行计算

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = Net(28 * 28, 300, 100, 10).to(device)

定义损失函数和优化器

criterion = nn.CrossEntropyLoss() # 交叉熵损失函数

optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9) # 使用随机梯度下降优化器

开始训练

losses = [] # 保存训练集的损失值 acces = [] # 保存训练集准确率 eval_losses = [] # 保存测试集损失值 eval_acces = [] # 保存测试集准确率 for epoch in range(num_epoches): train_loss = 0 train_acc = 0 model.train() # 将模型转化为训练模式 # 动态修改参数学习率（每五次） if epoch % 5 == 0: optimizer.param_groups[0]['lr'] *= 0.9 # 参数组 print("学习率：{:.6f}".format(optimizer.param_groups[0]['lr']))

# 训练过程
for img, label in train_loader:
    img = img.to(device)         # 将图像移动到GPU上
    label = label.to(device)        # 将标签移动到GPU上
    # 正向传播
    out = model(img)        # 将图像输入到神经网络中进行预测
    loss = criterion(out, label)        # 计算预测值和标签之间的损失
    # 反向传播
    optimizer.zero_grad()       # 梯度归零，清空上一次计算的梯度
    loss.backward()     # 反向传播计算梯度
    optimizer.step()        # 更新权重参数
    # 记录误差
    train_loss += loss.item()       # 记录训练集损失值
    # 计算分类的准确率
    _, pred = out.max(1)        # 找到预测的标签
    num_correct = (pred == label).sum().item()       # 计算预测正确的数量
    acc = num_correct / img.shape[0]         # 计算准确率
    train_acc += acc         # 计算训练集准确率

# 保存训练集损失和准确率
losses.append(train_loss / len(train_loader))        # 计算训练集损失平均值
acces.append(train_acc / len(train_loader))      # 计算训练集准确率平均值

# 在测试集上检验效果
model.eval()        # 将模型转化为评估模式
test_loss = 0        # 保存测试集的损失值
test_acc = 0        # 保存测试集准确率
with torch.no_grad():
    for img, label in test_loader:
        img = img.to(device)
        label = label.to(device)
        img = img.view(img.size(0), -1)
        out = model(img)
        loss = criterion(out, label)
        # 记录误差
        test_loss += loss.item()
        # 记录准确率
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        test_acc += num_correct / img.shape[0]

# 保存测试集损失和准确率
eval_losses.append(test_loss / len(test_loader))
eval_acces.append(test_acc / len(test_loader))
print('epoch: {}, Train Loss: {:.4f}, Train Acc: {:.4f}, Test Loss: {:.4f}, Test Acc: {:.4f}'
      .format(epoch, train_loss / len(train_loader), train_acc / len(train_loader),
              test_loss / len(test_loader), test_acc / len(test_loader)))