迁移学习 - NSP 案例 - 模型训练和评估

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


def train_model_nsp(train_dataset, my_tokenizer, my_bert_model, device, 
                    batch_size=8, num_epochs=3, learning_rate=1e-5):
    """训练模型：NSP 任务"""
    # 1. 创建数据加载器（方式三）
    train_dataloader = get_dataloader_nsp(train_dataset, my_tokenizer, device, batch_size)
    
    # 2. 创建模型并移动到设备
    my_model = AiModel(my_bert_model).to(device)
    
    # 3. 冻结 BERT 参数
    for param in my_model.bert.parameters():
        param.requires_grad = False
    
    # 4. 创建损失函数和优化器
    criterion = nn.CrossEntropyLoss(reduction='mean')
    optimizer = optim.Adam(my_model.parameters(), lr=learning_rate)
    
    # 5. 设置模型为训练模式
    my_model.train()
    
    # 6. 分批次训练
    for epoch in range(num_epochs):
        total_loss = 0
        for i, (inputs, labels) in enumerate(train_dataloader):
            # 前向传播
            outputs = my_model(**inputs)
            
            # 计算损失
            loss = criterion(outputs, labels)
            
            # 梯度清零
            optimizer.zero_grad()
            
            # 反向传播
            loss.backward()
            
            # 梯度更新
            optimizer.step()
            
            total_loss += loss.item()
            
            # 每 20 批打印一次
            if (i + 1) % 20 == 0:
                print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}], Loss: {loss.item():.4f}')
        
        print(f'Epoch {epoch+1} 完成, 平均损失: {total_loss/len(train_dataloader):.4f}')
    
    # 7. 保存模型（模型名改为 NSP）
    torch.save(my_model.state_dict(), './model_NSP.pth')
    print('模型已保存')


def evaluate_model_nsp(test_dataset, my_tokenizer, my_bert_model, device, 
                      batch_size=8, model_path='./model_NSP.pth'):
    """模型评估：NSP 任务"""
    # 1. 创建数据加载器（方式三，shuffle=False）
    test_dataloader = get_dataloader_nsp(test_dataset, my_tokenizer, device, batch_size, shuffle=False)
    
    # 2. 创建模型并加载训练好的模型
    my_model = AiModel(my_bert_model).to(device)
    my_model.load_state_dict(torch.load(model_path))
    
    # 3. 初始化评估参数
    correct = 0
    total = 0
    
    # 4. 设置模型为评估模式
    my_model.eval()
    
    # 5. 迭代预测
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(tqdm(test_dataloader, start=1)):
            # 前向传播
            outputs = my_model(**inputs)
            
            # 获取预测结果
            _, predicted = torch.max(outputs, 1)
            
            # 统计预测正确的样本数
            correct += (predicted == labels).sum().item()
            total += labels.size(0)
            
            # 打印第一批的预测结果
            if i == 1:
                print(f"\n第一批预测结果:")
                for j in range(min(8, len(predicted))):
                    # 解码原始文本（跳过特殊符号）
                    original_text = my_tokenizer.decode(
                        inputs['input_ids'][j], 
                        skip_special_tokens=True,
                        clean_up_tokenization_spaces=True
                    )
                    # 预测结果
                    predicted_label = predicted[j].item()
                    # 真实标签
                    true_label = labels[j].item()
                    
                    print(f"  样本 {j+1}:")
                    print(f"    原始文本: {original_text}")
                    print(f"    预测: {predicted_label} ({'是' if predicted_label == 1 else '不是'})")
                    print(f"    真实: {true_label} ({'是' if true_label == 1 else '不是'})")
    
    # 6. 计算准确率
    accuracy = correct / total
    print(f'\n准确率: {accuracy:.4f} ({accuracy*100:.2f}%)')
    
    return accuracy


# 测试
if __name__ == "__main__":
    print("=== 迁移学习：NSP 案例 - 模型训练和评估 ===")
    print("训练和评估函数已定义")
    print("准确率约 88%")