基于RNN的NLP学习（实战一）
代码来源：快速入门
代码所使用的使用的数据集：TorchVision
我的环境： python3.8，torch 2.2.2 ，torchdata 0.7.1，torchtext 0.17.2，torchvision 0.17.2
# 导入PyTorch及相关库
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
# 加载MNIST训练数据集，如果本地没有，则从网上下载
# MNIST是一个包含手写数字0-9的图像数据集，每个图像大小为28x28像素
training_data = datasets.MNIST(root='data', train=True, download=True, transform=ToTensor())
# 加载MNIST测试数据集，如果本地没有，则从网上下载
# 测试数据集包含用于评估模型性能的图像和标签
test_data = datasets.MNIST(root='data', train=False, download=True, transform=ToTensor())
# 设置批处理大小，即每次训练和测试时使用的样本数量
batch_size = 64
# 创建训练数据加载器，用于批量加载数据以进行训练
train_dataloader = DataLoader(training_data, batch_size=batch_size)
# 创建测试数据加载器，用于批量加载数据以进行测试
test_dataloader = DataLoader(test_data, batch_size=batch_size)
# 检查可用的设备，优先使用CUDA（GPU），其次是MPS（Apple的GPU），最后是CPU
# GPU加速可以显著提高训练效率
device = ("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
print(f"Using {device} device")
# 定义神经网络模型
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
# 展平层，将28x28的图像转换为784的向量
self.flatten = nn.Flatten()
# 线性层和ReLU激活函数的堆叠
self.linear_relu_stack = nn.Sequential(
nn.Linear(28 * 28, 512), # 第一层，输入784个节点，输出512个节点
nn.ReLU(), # ReLU激活函数，增加模型的非线性
nn.Linear(512, 512), # 第二层，输入512个节点，输出512个节点
nn.ReLU(), # ReLU激活函数
nn.Linear(512, 10), # 第三层，输入512个节点，输出10个节点（对应10个数字类别）
)
# 前向传播定义，指定数据如何通过模型
def forward(self, x):
x = self.flatten(x) # 展平输入图像，从[28, 28]到[784]
logits = self.linear_relu_stack(x) # 通过线性层和ReLU函数
return logits # 返回未归一化的概率（logits）
# 创建模型实例，并将其移动到定义的设备上
model = NeuralNetwork().to(device)
# 打印模型结构，以便了解模型的设计
print(model)
# 定义损失函数为交叉熵损失，用于分类问题
# 交叉熵损失比较模型输出的概率分布与实际标签的分布
loss_fn = nn.CrossEntropyLoss()
# 定义优化器为随机梯度下降，学习率为0.001
# 优化器负责更新模型参数以最小化损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
# 定义训练函数，用于训练模型
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset) # 训练集的大小
model.train() # 设置模型为训练模式，启用dropout和batch normalization
for batch, (X, y) in enumerate(dataloader): # 遍历数据加载器
X, y = X.to(device), y.to(device) # 将数据移动到定义的设备上
pred = model(X) # 计算模型预测
loss = loss_fn(pred, y) # 计算损失
loss.backward() # 反向传播计算梯度
optimizer.step() # 更新模型参数
optimizer.zero_grad() # 清空梯度，为下一次迭代准备
if batch % 100 == 0: # 每隔100个批次打印一次损失
loss, current = loss.item(), batch * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
# 定义测试函数，用于评估模型性能
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset) # 测试集的大小
num_batches = len(dataloader) # 测试批次的数量
model.eval() # 设置模型为评估模式，禁用dropout和batch normalization
# 初始化测试损失和正确数量
test_loss, correct = 0, 0
with torch.no_grad(): # 在不需要计算梯度的情况下
for X, y in dataloader: # 遍历数据加载器
X, y = X.to(device), y.to(device) # 将数据移动到定义的设备上
pred = model(X) # 计算模型预测
test_loss += loss_fn(pred, y).item() # 累加损失
correct += (pred.argmax(1) == y).type(torch.float).sum().item() # 累加正确预测的数量
test_loss /= num_batches # 计算平均损失
correct /= size # 计算准确率
# 打印测试错误和准确率
print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
# 设置训练的轮数
epochs = 5
for t in range(epochs): # 对每个轮数进行训练和测试
print(f"Epoch {t + 1}\n-------------------------------")
train(train_dataloader, model, loss_fn, optimizer) # 调用训练函数
test(test_dataloader, model, loss_fn) # 调用测试函数
print("Done!")
# 保存训练好的模型状态到文件
torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")
# 创建一个新的模型实例
model = NeuralNetwork().to(device)
# 从文件加载模型状态
model.load_state_dict(torch.load("model.pth"))
# 定义类别名称列表，用于打印预测结果
classes = [
"T-shirt/top",
"Trouser",
"Pullover",
"Dress",
"Coat",
"Sandal",
"Shirt",
"Sneaker",
"Bag",
"Ankle boot",
]
# 将模型设置为评估模式
model.eval()
# 获取测试数据集中的第一个图像和标签
x, y = test_data[0][0], test_data[0][1]
with torch.no_grad(): # 在不需要计算梯度的情况下
x = x.to(device) # 将图像移动到定义的设备上
pred = model(x) # 计算模型预测
# 获取预测和实际的类别名称
predicted, actual = classes[pred[0].argmax(0)], classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"') # 打印预测和实际结果