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()