train.py

import numpy as np
import torch
import os
import random
from model import ATModel
import argparse
from utils import *
import matplotlib.pyplot as plt
import pandas as pd
import matplotlib.ticker as ticker
import torch

def visualize_attention(attention_weights, tokens_p, tokens_q, title="Attention Heatmap"):
    """
    可视化注意力权重的热力图
    :param attention_weights: 注意力权重矩阵，形状为 [max_len, max_len]
    :param tokens_p: 句子P的词汇列表
    :param tokens_q: 句子Q的词汇列表
    :param title: 图的标题
    """
    plt.figure(figsize=(10, 8))
    ax = plt.gca()
    
    # 绘制热力图
    heatmap = ax.imshow(attention_weights.cpu().detach().numpy(), cmap="viridis", aspect="auto")
    
    # 设置坐标轴标签
    ax.set_xticks(range(len(tokens_q)))
    ax.set_yticks(range(len(tokens_p)))
    ax.set_xticklabels(tokens_q, rotation=90, fontsize=10)
    ax.set_yticklabels(tokens_p, fontsize=10)
    
    # 添加颜色条
    plt.colorbar(heatmap)
    
    # 设置标题
    plt.title(title, fontsize=14)
    plt.xlabel("Tokens in Sentence Q", fontsize=12)
    plt.ylabel("Tokens in Sentence P", fontsize=12)
    
    # 显示图像
    plt.tight_layout()
    plt.savefig(f"{title.replace(' ', '_')}.png")
    plt.show()



# 训练模型
def train(args, model, vocab, stop_words):
    model.train()
    train_data_loader, eval_x, eval_y = load_data(args.batch_size, vocab, stop_words, args)

    eval_p = eval_x[0]
    eval_q = eval_x[1]
    eval_p = torch.from_numpy(eval_p)
    eval_q = torch.from_numpy(eval_q)

    if torch.cuda.is_available():
        model = model.cuda()
        eval_p = eval_p.cuda()
        eval_q = eval_q.cuda()

    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    loss_func = nn.BCELoss()  # 二元交叉熵loss

    best_acc = 0

    for epoch in range(args.max_epoch):
        for step, (b_p, b_q, b_y) in enumerate(train_data_loader):
            # b_p.shape = b_q.shape = [batch_size, max_len]
            # b_y.shape = [batch_size]
            if torch.cuda.is_available():
                b_p = b_p.cuda()
                b_q = b_q.cuda()
                b_y = b_y.cuda()
            
            # TODO 1: 实现训练的前向传播和反向传播
            #前向传播
            output, cross_attention, self_attn_p, self_attn_q = model(b_p.long(), b_q.long())
            loss = loss_func(output.squeeze(), b_y.float())
            #反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()  
            
            if step % 100 == 0:
                # 每过完100个batch后，测试一下accuracy
                model.eval()
                with torch.no_grad():
                    test_output, cross_attention, self_attn_p, self_attn_q = model(eval_p.long(), eval_q.long())
                model.train()
                
                pred_y = (test_output.cpu().data.numpy() > 0.5).astype(int)
                accuracy = float((pred_y == eval_y).astype(int).sum()) / float(eval_y.size)
                if accuracy > best_acc:
                    best_acc = accuracy
                    torch.save({
                        'model_state_dict': model.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'accuracy': accuracy,
                        'epoch': epoch
                    }, 'ATmodel_best.pt')
                    print('save model, accuracy: %.3f' % accuracy)
                
                message = f'Epoch:{epoch} | Step: {step} | train loss: {loss.cpu().data.numpy():.4f} | test accuracy: {accuracy:.4f}'
                print(message)
                # 记录测试结果
                with open(args.log_file, 'a') as f:
                    f.write(f'{message}\n')

    # TODO 2: 实现注意力权重的可视化
    # 选择测试集中的样本，绘制自注意力和跨注意力的热力图
    # 需要显示词汇标签，设置合适的颜色映射和坐标轴
    
    vocab_inv = {idx: word for word, idx in vocab.items()}
    
    # 随机选择一个样本
    sample_idx = random.randint(0, eval_p.size(0) - 1)
    sample_p = eval_p[sample_idx]
    sample_q = eval_q[sample_idx]
    
    # 将样本转换为词汇列表
    tokens_p = [vocab_inv[idx] for idx in sample_p.cpu().numpy() if idx in vocab_inv]
    tokens_q = [vocab_inv[idx] for idx in sample_q.cpu().numpy() if idx in vocab_inv]
    
    # 获取注意力权重
    model.eval()
    with torch.no_grad():
        _, cross_attention, self_attn_p, self_attn_q = model(sample_p.unsqueeze(0).long(), sample_q.unsqueeze(0).long())
        
    # 可视化自注意力权重
    visualize_attention(self_attn_p[0], tokens_p, tokens_p, title="Self Attention Heatmap for Sentence P")
    visualize_attention(self_attn_q[0], tokens_q, tokens_q, title="Self Attention Heatmap for Sentence Q")
    # 可视化跨注意力权重
    visualize_attention(cross_attention[0], tokens_p, tokens_q, title="Cross Attention Heatmap")



if __name__ == '__main__':
    # 设置随机种子保证可重现性
    np.random.seed(42)
    torch.manual_seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(42)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
    # 配置参数
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--batch_size', type=int, default=128)
    argparser.add_argument('--del_word_frequency', type=int, default=3)
    argparser.add_argument('--char_dim', type=int, default=200)
    argparser.add_argument('--char_hidden_size', type=int, default=128)
    argparser.add_argument('--max_len', type=int, default=12)
    argparser.add_argument('--lr', type=float, default=0.001)
    argparser.add_argument('--max_epoch', type=int, default=1)
    argparser.add_argument('--num_layers', type=int, default=2, help='Transformer层数')
    argparser.add_argument('--num_heads', type=int, default=8, help='注意力头数')
    argparser.add_argument('--ff_hidden_ratio', type=int, default=4, help='FFN隐藏层维度比例')

    argparser.add_argument('--log_file_path', type=str, default='./log')
    args = argparser.parse_args()
    
    # 配置日志文件,记录训练过程
    cur_time = get_local_time()
    args_str = f'week6-{cur_time}'
    args.log_file = os.path.join(args.log_file_path, args_str + '.txt')
    print(str(args))
    with open(args.log_file, 'a') as f:
        f.write(str(args) + '\n')

    # ---1. 加载停用词、构建词典、读取数据---
    stop_words = load_stop_words()
    vocab = {}
    if not os.path.exists('./data/vocab.txt'):
        build_vocab(del_word_frequency=args.del_word_frequency, stop_words=stop_words)

    with open('./data/vocab.txt', encoding='utf-8') as file:
        for line in file.readlines():
            vocab[line.strip()] = len(vocab)

    # ---2. 初始化模型---
    vocab_size = len(vocab)
    model = ATModel(
        vocab_size=vocab_size,
        char_dim=args.char_dim,
        char_hidden_size=args.char_hidden_size,
        max_len=args.max_len,
        num_heads=args.num_heads,
        num_layers=args.num_layers,
        ff_hidden_ratio=args.ff_hidden_ratio
    )
    
    print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")

    # ---3. 开始训练---
    train(args, model, vocab, stop_words)