code/transformercode/train.py

from config import *
from utils import *
from model import make_model
from data_loader import *
from torch.nn import CrossEntropyLoss
# from torch.optim import Adam

from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR

import time
from datetime import datetime
from pathlib import Path
from torch.utils.tensorboard import SummaryWriter

#自定义调整策略
def lr_lambda_fn(step, wramup):
    lr = 0
    if step <= wramup:
        lr = step / wramup * 10
    else:
        lr = wramup / step * 10
    return max(lr, 0.1)

def run_epoch(loader, model, loss_fn, optimizer=None):
    # 初始化loss值，和batch数量总数
    total_batchs = 0.
    total_loss = 0.
    model.to(device)
    # 加载数据并开始训练
    for src_x, src_mask, tgt_x, tgt_mask, tgt_y, tgt_text in loader:
        src_x= src_x.to(device)
        src_mask=src_mask.to(device)
        tgt_x=tgt_x.to(device)
        tgt_mask=tgt_mask.to(device)
        tgt_y=tgt_y.to(device)#9,6
        output = model(src_x, src_mask, tgt_x, tgt_mask)#9,6,19
        # 交叉熵损失，要求目标值是一维的
        # out.shape 54,19     label.shape 54
        loss = loss_fn(output.reshape(-1, output.shape[-1]), tgt_y.reshape(-1))
        # 累积batch数量和loss值
        total_batchs += 1
        total_loss += loss.item()

        # 如果有优化器，则表示需要反向传播
        if optimizer:
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    # 返回epoch的平均loss值
    return total_loss / total_batchs

#评估函数
def evaluate(loader, model, max_len=MAX_LEN):
    tgt_sent = [] #label 标准中文答案
    prob_sent = [] #pre 模型生成的中文

    for src_x, src_mask, tgt_x, tgt_mask, tgt_y, tgt_text in loader:
        batch_prob_text = batch_greedy_decode(model, src_x, src_mask, max_len)
        tgt_sent += tgt_text
        prob_sent += batch_prob_text

    # 验证集包含数万条句子，完整打印会淹没终端并拖慢训练。
    # 只展示少量预测/参考样例，完整指标仍使用全部验证集计算。
    preview_size = min(3, len(prob_sent))
    for index in range(preview_size):
        print(f"验证样例 {index + 1} 预测: {prob_sent[index]}")
        print(f"验证样例 {index + 1} 参考: {tgt_sent[index]}")
    return bleu_score(prob_sent, [tgt_sent])

if __name__== "__main__":
    # 创建日志目录，并为本次训练创建独立的日志文件和 TensorBoard 目录。
    run_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_dir = Path("./logs")
    log_dir.mkdir(parents=True, exist_ok=True)
    params_dir = Path("./params")
    params_dir.mkdir(parents=True, exist_ok=True)
    log_path = log_dir / f"train_{run_time}.log"
    tensorboard_dir = log_dir / f"tensorboard_{run_time}"
    writer = SummaryWriter(log_dir=str(tensorboard_dir))
    log_file = log_path.open("w", encoding="utf-8")

    # 记录训练开始时间和日志保存位置。
    training_start_time = time.time()
    start_message = (
        f"训练开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        f"日志文件: {log_path}\n"
        f"TensorBoard目录: {tensorboard_dir}"
    )
    print(start_message)
    log_file.write(start_message + "\n")
    log_file.flush()

    en_id2vocab, _  = get_vocab('en')
    zh_id2vocab, _ = get_vocab('zh')

    SRC_VOCAB_SIZE = len(en_id2vocab)
    TGT_VOCAB_SIZE = len(zh_id2vocab)

    model = make_model(SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, D_MODEL, N_HEAD, D_FF, N, DROPOUT)
    
    train_dataset = Dataset('train')
    train_loader = data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True,
                                   collate_fn=train_dataset.collate_fn)

    val_dataset = Dataset('val')
    val_loader = data.DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, 
                                 collate_fn=val_dataset.collate_fn)
    
    loss_fn = CrossEntropyLoss(ignore_index=PAD_ID,label_smoothing=LABEL_SMOOTHING)
    # optimizer = Adam(model.parameters(), lr=LR)

    optimizer = AdamW(model.parameters(), lr=LR)

#     optimizer = torch.optim.AdamW(
#     model.parameters(),
#     lr=LR,
#     weight_decay=0.01,
# )

    lr_scheduler = LambdaLR(optimizer, lr_lambda=lambda step: lr_lambda_fn(step, EPOCH/4))
    
    # 早停改进：以验证集BLEU为标准保存最佳权重。
    # 连续early_stop_patience轮没有至少min_bleu_delta的提升时停止。
    best_bleu = float("-inf")
    best_epoch = 0
    early_stop_patience = 6
    epochs_without_improvement = 0
    min_bleu_delta = 0.01

    for e in range(EPOCH):
        
        epoch_start_time = time.time()
        # 训练流程
        model.train()
        train_loss = run_epoch(train_loader, model, loss_fn, optimizer)
        lr_scheduler.step()
        # 打印当前学习率
        # current_lr = optimizer.param_groups[0]['lr']
        # print(current_lr)
        # 验证流程
        model.eval()
        val_loss = run_epoch(val_loader, model, loss_fn, None)
        #writer.add_scalars("loss",{'train_loss':train_loss, 'val_loss':val_loss},e)
        # print('epoch:', e, 'train_loss:', train_loss, 'val_loss:',val_loss)
        
        #评估模型
        val_bleu = evaluate(val_loader, model, MAX_LEN)
        # 计算单轮耗时，并将 loss、BLEU 和耗时同时打印到终端和日志文件。
        epoch_time = time.time() - epoch_start_time
        log_message = (
            f"epoch: {e + 1}/{EPOCH} "
            f"train_loss: {train_loss:.6f} "
            f"val_loss: {val_loss:.6f} "
            f"val_bleu: {val_bleu:.2f} "
            f"epoch_time: {epoch_time:.2f}s"
        )
        print(log_message)
        log_file.write(log_message + "\n")
        log_file.flush()

        # 分别记录训练 loss 和验证 loss，供 TensorBoard 绘制曲线。
        writer.add_scalar("Loss/train", train_loss, e + 1)
        writer.add_scalar("Loss/validation", val_loss, e + 1)
        writer.add_scalar("BLEU/validation", val_bleu, e + 1)
        writer.add_scalar("Time/epoch_seconds", epoch_time, e + 1)
        writer.flush()

        # 每2轮保存一次模型权重。e从0开始，因此使用e+1表示实际轮数。
        if (e + 1) % 2 == 0:
            epoch_model_path = params_dir / f"epoch_{e + 1}.pt"
            torch.save(model.state_dict(), epoch_model_path)
            print("定期模型已保存:", epoch_model_path)

        # 早停改进：最佳权重保存和未改善轮数统计必须使用同一套判断。
        # val_bleu已经保留两位小数，因此提升0.01及以上视为有效改善。
        if val_bleu >= best_bleu + min_bleu_delta:
            torch.save(
                model.state_dict(),
                params_dir / "best.pt",
            )

            best_bleu = val_bleu
            best_epoch = e + 1
            epochs_without_improvement = 0

            early_stop_status = (
                f"发现更好的BLEU：{best_bleu:.2f}，"
                f"已更新best.pt（epoch {best_epoch}）"
            )

        else:
            epochs_without_improvement += 1

            early_stop_status = (
                f"BLEU连续未改善轮数："
                f"{epochs_without_improvement}/"
                f"{early_stop_patience}，"
                f"当前最佳：{best_bleu:.2f}"
            )

        print(early_stop_status)
        log_file.write(early_stop_status + "\n")
        log_file.flush()
        writer.add_scalar(
            "EarlyStopping/epochs_without_improvement",
            epochs_without_improvement,
            e + 1,
        )
        writer.flush()

        # 连续多轮没有改善，提前结束
        if epochs_without_improvement >= early_stop_patience:
            early_stop_message = (
                f"触发早停：BLEU连续"
                f"{early_stop_patience}轮没有改善；"
                f"最佳BLEU为{best_bleu:.2f}，"
                f"出现在第{best_epoch}轮。"
            )
            print(early_stop_message)
            log_file.write(early_stop_message + "\n")
            log_file.flush()
            break

    # 训练结束后记录总耗时，并关闭日志文件和 TensorBoard 写入器。
    total_training_time = time.time() - training_start_time
    end_message = (
        f"训练结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} "
        f"总训练时间: {total_training_time:.2f}s "
        f"({total_training_time / 3600:.2f}小时)"
    )
    print(end_message)
    log_file.write(end_message + "\n")
    log_file.close()
    writer.close()