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

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

my_lr = 1e-3
epochs = 1

def dm_train_lstm():
    # 1. 数据准备（与 RNN 一致）
    from name_classification import read_data, NameClassDataset
    from torch.utils.data import DataLoader
    
    my_list_x, my_list_y = read_data()
    name_class_dataset = NameClassDataset(my_list_x, my_list_y)
    my_data_loader = DataLoader(name_class_dataset, batch_size=1, shuffle=True)
    
    # 2. 模型与优化器初始化
    input_size = 57
    hidden_size = 128
    output_size = 18
    
    # 修改点1：模型实例化（变量名建议不改）
    my_lstm = MyLSTM(input_size, hidden_size, output_size, n_layers=1)
    
    criterion = nn.NLLLoss()
    optimizer = optim.Adam(my_lstm.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):
            # 修改点2：initHidden 返回 hidden 和 c
            hidden, c = my_lstm.initHidden()
            
            # 修改点3：前向传播传入 hidden 和 c，返回 output, hn, cn
            output, hn, cn = my_lstm(x, hidden, c)
            
            # 计算损失
            loss = criterion(output, y)
            
            # 三剑客
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            current_loss += loss.item()
            step += 1
            
            pred = output.argmax(dim=1)
            correct += (pred == y).sum().item()
            
            if step % 100 == 0:
                avg_loss = current_loss / 100
                loss_list.append(avg_loss)
                current_loss = 0
            
            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
    
    end_time = time.time()
    print(f"训练总时间: {end_time - start_time:.2f} 秒")
    
    # 修改点4：保存模型
    torch.save(my_lstm.state_dict(), 'my_lstm.pth')
    
    return loss_list, acc_list


loss_list, acc_list = dm_train_lstm()