【“Transformers快速入门”学习笔记6】训练模型，保存和加载模型

Pytorch 所有的模块（层）都是 nn.Module 的子类，神经网络模型本身就是一个模块，它还包含了很多其他的模块。
下面给出了一个完整的训练循环和测试循环过程，以及详细注释

训练模型

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# 检测设备是否可用 CUDA，若可用则使用 CUDA 加速，否则使用 CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using {device} device')

# 下载并加载 FashionMNIST 数据集作为训练集和测试集
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()  # 转换图像为 PyTorch 张量
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()  # 转换图像为 PyTorch 张量
)

learning_rate = 1e-3  # 设置学习率
batch_size = 64       # 设置每个批次的样本数
epochs = 3            # 设置训练循环的轮数

# 创建训练集和测试集的 DataLoader 对象
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

# 定义神经网络模型类
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()  # 将输入的二维图像展平为一维
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),  # 全连接层，输入维度为 28*28，输出维度为 512
            nn.ReLU(),               # ReLU 激活函数
            nn.Linear(512, 256),     # 全连接层，输入维度为 512，输出维度为 256
            nn.ReLU(),               # ReLU 激活函数
            nn.Linear(256, 10),      # 全连接层，输入维度为 256，输出维度为 10（对应分类数）
            nn.Dropout(p=0.2)        # Dropout 层，防止过拟合
        )

    def forward(self, x):
        x = self.flatten(x)               # 将输入展平
        logits = self.linear_relu_stack(x) # 前向传播计算预测值
        return logits

model = NeuralNetwork().to(device)  # 实例化神经网络模型并将其移动到设备上（CPU 或 GPU）

# 定义训练函数，用于训练神经网络模型
def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()  # 将模型设为训练模式
    for batch, (X, y) in enumerate(dataloader, start=1):
        X, y = X.to(device), y.to(device)  # 将输入数据移动到设备上
        # 计算预测值和损失
        pred = model(X)
        loss = loss_fn(pred, y)
        # 反向传播
        optimizer.zero_grad()  # 梯度清零
        loss.backward()        # 反向传播计算梯度
        optimizer.step()       # 更新模型参数

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

# 定义测试函数，用于评估神经网络模型在测试集上的性能
def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    model.eval()  # 将模型设为评估模式
    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(dim=-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")

loss_fn = nn.CrossEntropyLoss()  # 定义损失函数为交叉熵损失函数
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)  # 定义优化器为 AdamW

# 训练循环
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)  # 训练模型
    test_loop(test_dataloader, model, loss_fn)  # 在测试集上评估模型性能
print("Done!")  # 训练结束提示

注意：一定要在预测之前调用 model.eval() 方法将 dropout 层和 batch normalization 层设置为评估模式，否则会产生不一致的预测结果。

优化模型参数

模型训练是一个迭代的过程，每一轮 epoch 迭代中模型都会对输入样本进行预测，然后对预测结果计算损失 (loss)，并求 loss 对每一个模型参数的偏导，最后使用优化器更新所有的模型参数。

损失函数 (Loss function) 用于度量预测值与答案之间的差异，模型的训练过程就是最小化损失函数。Pytorch 实现了很多常见的损失函数，例如用于回归任务的均方误差 (Mean Square Error) nn.MSELoss、用于分类任务的负对数似然 (Negative Log Likelihood) nn.NLLLoss、同时结合了 nn.LogSoftmax 和 nn.NLLLoss 的交叉熵损失 (Cross Entropy) nn.CrossEntropyLoss 等。
优化器 (Optimization) 使用特定的优化算法（例如随机梯度下降），通过在每一个训练阶段 (step) 减少（基于一个 batch 样本计算的）模型损失来调整模型参数。Pytorch 实现了很多优化器，例如 SGD、ADAM、RMSProp 等。

每一轮迭代 (Epoch) 实际上包含了两个步骤：

训练循环 (The Train Loop) 在训练集上进行迭代，尝试收敛到最佳的参数；
验证/测试循环 (The Validation/Test Loop) 在测试/验证集上进行迭代以检查模型性能有没有提升。
具体地，在训练循环中，优化器通过以下三个步骤进行优化：
调用 optimizer.zero_grad() 重设模型参数的梯度。默认情况下梯度会进行累加，为了防止重复计算，在每个训练阶段开始前都需要清零梯度；
通过 loss.backwards() 反向传播预测结果的损失，即计算损失对每一个参数的偏导；
调用 optimizer.step() 根据梯度调整模型的参数。

保存及加载模型

Pytorch 模型会将所有参数存储在一个状态字典 (state dictionary) 中。
通过**model.state_dict()**加载，通过torch.save()保存：

保存和加载模型权重

import torch
import torchvision.models as models

model = models.vgg16(pretrained=True)
torch.save(model.state_dict(), 'model_weights.pth')# 保存
model.load_state_dict(torch.load('model_weights.pth'))# 加载
model.eval()

保存和加载完整模型

import torch
import torchvision.models as models

model = models.vgg16(pretrained=True)
torch.save(model, 'model.pth')
model = torch.load('model.pth')