全球人名分类案例 - RNN 模型训练

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

# 1. 定义变量，记录学习率和训练轮数
my_lr = 1e-3      # 学习率 0.001
epochs = 1        # 训练轮数

def dm_train_rnn():
    # ========== 1. 数据准备 ==========
    # 1.1 读取数据
    from name_classification import read_data
    my_list_x, my_list_y = read_data()
    
    # 1.2 构建数据集对象
    from name_classification import NameClassDataset
    name_class_dataset = NameClassDataset(my_list_x, my_list_y)
    
    # 1.3 构建数据加载器
    from torch.utils.data import DataLoader
    my_data_loader = DataLoader(
        name_class_dataset,
        batch_size=1,
        shuffle=True
    )
    
    # ========== 2. 模型与优化器初始化 ==========
    # 2.1 定义模型参数
    input_size = 57       # 字符表大小
    hidden_size = 128     # 隐藏层维度
    output_size = 18      # 国家数量
    
    # 2.2 创建模型对象
    my_rnn = MyRNN(input_size, hidden_size, output_size, n_layers=1)
    
    # 2.3 定义损失函数和优化器
    # 如果用了 LogSoftmax，则用 NLLLoss
    criterion = nn.NLLLoss()
    # 如果用 CrossEntropyLoss，则模型不需要 LogSoftmax
    # criterion = nn.CrossEntropyLoss()
    
    optimizer = optim.Adam(my_rnn.parameters(), lr=my_lr)
    
    # ========== 3. 参数初始化 ==========
    import time
    start_time = time.time()  # 记录开始时间
    loss_list = []             # 损失列表（用于绘图）
    acc_list = []              # 准确率列表（用于绘图）
    
    # ========== 4. 训练过程 ==========
    for epoch in range(epochs):
        # 当前批次累计损失
        current_loss = 0
        # 训练步数
        step = 0
        # 正确预测数
        correct = 0
        
        # 内循环：每批次遍历
        for x, y in tqdm(my_data_loader):
            # 4.1 获取结果
            # x 形状：(seq_len, 1, 57)
            # y 形状：(1,)
            
            # 4.2 初始化隐藏状态
            hidden = my_rnn.initHidden()
            
            # 4.3 前向传播
            output, hn = my_rnn(x, hidden)
            
            # 4.4 计算损失
            loss = criterion(output, y)
            
            # 4.5 三剑客：梯度清零、反向传播、参数更新
            optimizer.zero_grad()  # 梯度清零
            loss.backward()        # 反向传播
            optimizer.step()       # 参数更新
            
            # 4.6 累加损失
            current_loss += loss.item()
            step += 1
            
            # 4.7 预测对错
            pred = output.argmax(dim=1)
            correct += (pred == y).sum().item()
            
            # 4.8 每 100 次求平均损失，存入列表
            if step % 100 == 0:
                avg_loss = current_loss / 100
                loss_list.append(avg_loss)
                current_loss = 0
            
            # 4.9 每 2000 次打印日志
            if step % 2000 == 0:
                avg_acc = correct / 2000
                acc_list.append(avg_acc)
                print(f"轮数: {epoch+1}, 步数: {step}, "
                      f"损失: {avg_loss:.4f}, 正确率: {avg_acc:.4f}")
                correct = 0
    
    # ========== 5. 训练结束 ==========
    end_time = time.time()
    print(f"训练总时间: {end_time - start_time:.2f} 秒")
    
    return loss_list, acc_list


# 运行训练
loss_list, acc_list = dm_train_rnn()