410 lines
17 KiB
Python
410 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
临时迁移训练脚本：在预训练模型基础上重新训练，支持冻结 context_encoder。
|
|
|
|
与 src/model/trainer.py 的 train 命令行为完全一致，额外增加：
|
|
- --pretrained-checkpoint: 加载预训练权重（必需的迁移学习源）
|
|
- --freeze-context-encoder: 冻结 context_encoder 层（默认开启）
|
|
|
|
运行方式:
|
|
python scripts/finetune_slots.py \
|
|
--pretrained-checkpoint ./output/checkpoints/best_model.pt \
|
|
--train-data-path /path/to/train_data \
|
|
--eval-data-path /path/to/eval_data \
|
|
--output-dir ./finetune_output \
|
|
--freeze-context-encoder
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
import sys
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))
|
|
|
|
from model.model import InputMethodEngine
|
|
from model.trainer import (
|
|
Trainer,
|
|
create_dataloader,
|
|
worker_init_fn,
|
|
)
|
|
from model.dataset import PinyinInputDataset
|
|
from model.preprocessed_dataset import (
|
|
PreProcessedDataset,
|
|
is_preprocessed_data,
|
|
)
|
|
|
|
from loguru import logger
|
|
from rich.console import Console
|
|
from rich.panel import Panel
|
|
from rich.table import Table
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="迁移学习训练：加载预训练模型，冻结指定层后重新训练",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
|
|
# === 数据参数 ===
|
|
parser.add_argument("--train-data-path", "-t", required=True, help="训练数据集路径")
|
|
parser.add_argument("--eval-data-path", "-e", required=True, help="评估数据集路径")
|
|
parser.add_argument("--output-dir", "-o", default="./finetune_output", help="输出目录")
|
|
parser.add_argument("--max-iter-length", type=int,
|
|
default=1024 * 1024 * 128, help="每个 epoch 最大样本数")
|
|
|
|
# === 迁移学习参数 ===
|
|
parser.add_argument("--pretrained-checkpoint", "-c", required=True,
|
|
help="预训练模型检查点路径")
|
|
parser.add_argument("--freeze-context-encoder", action="store_true", default=True,
|
|
help="冻结 context_encoder 层 (默认开启)")
|
|
parser.add_argument("--no-freeze-context-encoder", dest="freeze_context_encoder",
|
|
action="store_false",
|
|
help="不冻结 context_encoder")
|
|
|
|
# === 训练参数 ===
|
|
parser.add_argument("--batch-size", "-b", type=int, default=128, help="批次大小")
|
|
parser.add_argument("--num-epochs", type=int, default=10, help="训练轮数")
|
|
parser.add_argument("--learning-rate", "-lr", type=float, default=2e-4,
|
|
help="学习率")
|
|
parser.add_argument("--min-learning-rate", type=float, default=1e-9,
|
|
help="最小学习率")
|
|
parser.add_argument("--weight-decay", type=float, default=0.05, help="权重衰减")
|
|
parser.add_argument("--warmup-ratio", type=float, default=0.1, help="热身步数比例")
|
|
parser.add_argument("--label-smoothing", type=float, default=0.1,
|
|
help="标签平滑参数")
|
|
parser.add_argument("--grad-accum-steps", type=int, default=1,
|
|
help="梯度累积步数")
|
|
parser.add_argument("--clip-grad-norm", type=float, default=1.0,
|
|
help="梯度裁剪范数")
|
|
parser.add_argument("--eval-frequency", type=int, default=500, help="评估频率")
|
|
parser.add_argument("--save-frequency", type=int, default=1000, help="保存频率")
|
|
|
|
# === 其他参数 ===
|
|
parser.add_argument("--mixed-precision", action="store_true", default=True)
|
|
parser.add_argument("--no-mixed-precision", dest="mixed_precision",
|
|
action="store_false", help="禁用混合精度")
|
|
parser.add_argument("--num-workers", type=int, default=2, help="数据加载worker数")
|
|
parser.add_argument("--tensorboard", action="store_true", default=True)
|
|
parser.add_argument("--no-tensorboard", dest="tensorboard", action="store_false",
|
|
help="禁用 TensorBoard")
|
|
parser.add_argument("--seed", type=int, default=42, help="随机种子")
|
|
parser.add_argument("--compile", action="store_true", default=False,
|
|
help="使用 torch.compile 优化")
|
|
parser.add_argument("--moe-mode", default="all",
|
|
choices=["all", "sparse", "sparse_allow_graph"],
|
|
help="MoE 计算策略")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# ================================================================
|
|
# 初始化
|
|
# ================================================================
|
|
torch.multiprocessing.set_sharing_strategy("file_system")
|
|
if torch.cuda.is_available():
|
|
torch.set_float32_matmul_precision("high")
|
|
|
|
torch.manual_seed(args.seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(args.seed)
|
|
|
|
console = Console()
|
|
output_path = Path(args.output_dir)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# ================================================================
|
|
# 模型常量 (与 trainer.py 保持一致)
|
|
# ================================================================
|
|
vocab_size = 10019
|
|
pinyin_vocab_size = 30
|
|
dim = 512
|
|
num_slots = 8
|
|
n_layers = 4
|
|
n_heads = 4
|
|
num_experts = 10
|
|
max_seq_len = 128
|
|
|
|
# ================================================================
|
|
# 打印配置
|
|
# ================================================================
|
|
console.print(Panel.fit(
|
|
"[bold cyan]迁移学习训练配置[/bold cyan]", border_style="cyan"))
|
|
config_table = Table(show_header=True, header_style="bold magenta")
|
|
config_table.add_column("Category", style="cyan")
|
|
config_table.add_column("Parameter", style="green")
|
|
config_table.add_column("Value", style="yellow")
|
|
|
|
config_table.add_row("迁移学习", "预训练检查点", args.pretrained_checkpoint)
|
|
config_table.add_row("迁移学习", "冻结 context_encoder",
|
|
str(args.freeze_context_encoder))
|
|
config_table.add_row("数据", "训练数据路径", args.train_data_path)
|
|
config_table.add_row("数据", "评估数据路径", args.eval_data_path)
|
|
config_table.add_row("数据", "输出目录", args.output_dir)
|
|
config_table.add_row("数据", "批次大小", str(args.batch_size))
|
|
config_table.add_row("数据", "Worker数量", str(args.num_workers))
|
|
config_table.add_row("模型", "MoE策略", args.moe_mode)
|
|
config_table.add_row("模型", "编译优化", str(args.compile))
|
|
config_table.add_row("训练", "训练轮数", str(args.num_epochs))
|
|
config_table.add_row("训练", "学习率", f"{args.learning_rate:.2e}")
|
|
config_table.add_row("训练", "最小学习率", f"{args.min_learning_rate:.2e}")
|
|
config_table.add_row("训练", "权重衰减", str(args.weight_decay))
|
|
config_table.add_row("训练", "热身比例", str(args.warmup_ratio))
|
|
config_table.add_row("训练", "标签平滑", str(args.label_smoothing))
|
|
config_table.add_row("训练", "梯度累积", str(args.grad_accum_steps))
|
|
config_table.add_row("训练", "梯度裁剪", str(args.clip_grad_norm))
|
|
config_table.add_row("训练", "混合精度", str(args.mixed_precision))
|
|
|
|
# ================================================================
|
|
# 创建数据加载器 (逻辑与 trainer.py CLI 完全一致)
|
|
# ================================================================
|
|
console.print("[bold cyan]正在创建数据加载器...[/bold cyan]")
|
|
|
|
is_train_preprocessed = is_preprocessed_data(args.train_data_path)
|
|
is_eval_preprocessed = is_preprocessed_data(args.eval_data_path)
|
|
|
|
if is_train_preprocessed:
|
|
train_dataset = PreProcessedDataset(args.train_data_path,
|
|
max_cache_shards=2)
|
|
pre_shuffled = train_dataset.metadata.get("pre_shuffled", False)
|
|
shuffle_train = not pre_shuffled
|
|
if args.max_iter_length > 0:
|
|
capped_samples = min(len(train_dataset), args.max_iter_length)
|
|
else:
|
|
capped_samples = len(train_dataset)
|
|
total_steps = (capped_samples // args.batch_size) * args.num_epochs
|
|
train_num_workers = min(args.num_workers, 1)
|
|
logger.info(
|
|
f"Preprocessed dataset: {len(train_dataset):,} samples, "
|
|
f"shuffle={shuffle_train}, pre_shuffled={pre_shuffled}, "
|
|
f"workers={train_num_workers}, steps={total_steps:,}")
|
|
train_dataloader = create_dataloader(
|
|
dataset=train_dataset,
|
|
batch_size=args.batch_size,
|
|
num_workers=train_num_workers,
|
|
pin_memory=torch.cuda.is_available(),
|
|
shuffle=shuffle_train,
|
|
)
|
|
config_table.add_row("数据", "训练数据类型", "预处理数据")
|
|
else:
|
|
train_dataset = PinyinInputDataset(
|
|
data_path=args.train_data_path,
|
|
max_workers=-1,
|
|
max_iter_length=args.max_iter_length,
|
|
max_seq_length=max_seq_len,
|
|
text_field="text",
|
|
py_style_weight=(9, 2, 1),
|
|
shuffle_buffer_size=2000000,
|
|
length_weights={1: 10, 2: 50, 3: 50, 4: 40,
|
|
5: 15, 6: 10, 7: 5, 8: 2},
|
|
)
|
|
total_steps = int(args.max_iter_length *
|
|
args.num_epochs / args.batch_size)
|
|
train_dataloader = create_dataloader(
|
|
dataset=train_dataset,
|
|
batch_size=args.batch_size,
|
|
num_workers=args.num_workers,
|
|
pin_memory=torch.cuda.is_available(),
|
|
max_iter_length=args.max_iter_length,
|
|
)
|
|
config_table.add_row("数据", "训练数据类型", "流式数据")
|
|
|
|
if is_eval_preprocessed:
|
|
eval_dataset = PreProcessedDataset(args.eval_data_path,
|
|
max_cache_shards=1)
|
|
eval_dataloader = create_dataloader(
|
|
dataset=eval_dataset,
|
|
batch_size=args.batch_size,
|
|
num_workers=0,
|
|
pin_memory=torch.cuda.is_available(),
|
|
shuffle=False,
|
|
)
|
|
config_table.add_row("数据", "评估数据类型", "预处理数据")
|
|
else:
|
|
eval_dataset = PinyinInputDataset(
|
|
data_path=args.eval_data_path,
|
|
max_workers=-1,
|
|
max_iter_length=args.batch_size * 64,
|
|
max_seq_length=max_seq_len,
|
|
text_field="text",
|
|
py_style_weight=(9, 2, 1),
|
|
shuffle_buffer_size=2000000,
|
|
length_weights={1: 10, 2: 50, 3: 50, 4: 40,
|
|
5: 15, 6: 10, 7: 5, 8: 2},
|
|
)
|
|
eval_dataloader = create_dataloader(
|
|
dataset=eval_dataset,
|
|
batch_size=args.batch_size,
|
|
num_workers=2,
|
|
pin_memory=torch.cuda.is_available(),
|
|
max_iter_length=args.batch_size * 64,
|
|
)
|
|
config_table.add_row("数据", "评估数据类型", "流式数据")
|
|
|
|
config_table.add_row("数据", "总步数", str(total_steps))
|
|
console.print(config_table)
|
|
|
|
# ================================================================
|
|
# 创建模型并加载预训练权重
|
|
# ================================================================
|
|
console.print("[bold cyan]正在创建模型并加载预训练权重...[/bold cyan]")
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = InputMethodEngine(
|
|
vocab_size=vocab_size,
|
|
pinyin_vocab_size=pinyin_vocab_size,
|
|
dim=dim,
|
|
num_slots=num_slots,
|
|
n_layers=n_layers,
|
|
n_heads=n_heads,
|
|
num_experts=num_experts,
|
|
max_seq_len=max_seq_len,
|
|
compile=args.compile,
|
|
moe_mode=args.moe_mode,
|
|
)
|
|
model.to(device)
|
|
|
|
# 加载预训练权重
|
|
pretrained_path = Path(args.pretrained_checkpoint)
|
|
if not pretrained_path.exists():
|
|
console.print(f"[red]❌ 预训练检查点不存在: {args.pretrained_checkpoint}[/red]")
|
|
sys.exit(1)
|
|
|
|
checkpoint = torch.load(args.pretrained_checkpoint, map_location=device)
|
|
if "model_state_dict" in checkpoint:
|
|
pretrained_weights = checkpoint["model_state_dict"]
|
|
else:
|
|
pretrained_weights = checkpoint
|
|
|
|
missing_keys, unexpected_keys = model.load_state_dict(
|
|
pretrained_weights, strict=False)
|
|
|
|
if missing_keys:
|
|
console.print(f"[yellow]⚠ 缺失的键 ({len(missing_keys)}): "
|
|
f"{missing_keys[:5]}...[/yellow]")
|
|
if unexpected_keys:
|
|
console.print(f"[yellow]⚠ 多余的键 ({len(unexpected_keys)}): "
|
|
f"{unexpected_keys[:5]}...[/yellow]")
|
|
|
|
console.print(f"[green]✓ 预训练权重加载完成[/green]")
|
|
|
|
# ================================================================
|
|
# 冻结 context_encoder
|
|
# ================================================================
|
|
if args.freeze_context_encoder:
|
|
console.print("[bold cyan]正在冻结 context_encoder...[/bold cyan]")
|
|
frozen_count = 0
|
|
trainable_count = 0
|
|
for name, param in model.named_parameters():
|
|
if name.startswith("context_encoder"):
|
|
param.requires_grad = False
|
|
frozen_count += param.numel()
|
|
else:
|
|
trainable_count += param.numel()
|
|
|
|
total_params = frozen_count + trainable_count
|
|
console.print(f"[green]✓ context_encoder 已冻结[/green]")
|
|
logger.info(
|
|
f"冻结参数: {frozen_count:,} / {total_params:,} "
|
|
f"({frozen_count / total_params * 100:.1f}%), "
|
|
f"可训练参数: {trainable_count:,} / {total_params:,} "
|
|
f"({trainable_count / total_params * 100:.1f}%)")
|
|
else:
|
|
logger.info("未冻结任何层，全模型参与训练")
|
|
|
|
# ================================================================
|
|
# 保存配置
|
|
# ================================================================
|
|
config = {
|
|
"pretrained_checkpoint": args.pretrained_checkpoint,
|
|
"freeze_context_encoder": args.freeze_context_encoder,
|
|
"train_data_path": args.train_data_path,
|
|
"eval_data_path": args.eval_data_path,
|
|
"output_dir": args.output_dir,
|
|
"batch_size": args.batch_size,
|
|
"num_epochs": args.num_epochs,
|
|
"learning_rate": args.learning_rate,
|
|
"min_learning_rate": args.min_learning_rate,
|
|
"weight_decay": args.weight_decay,
|
|
"warmup_ratio": args.warmup_ratio,
|
|
"label_smoothing": args.label_smoothing,
|
|
"grad_accum_steps": args.grad_accum_steps,
|
|
"clip_grad_norm": args.clip_grad_norm,
|
|
"eval_frequency": args.eval_frequency,
|
|
"save_frequency": args.save_frequency,
|
|
"mixed_precision": args.mixed_precision,
|
|
"num_workers": args.num_workers,
|
|
"use_tensorboard": args.tensorboard,
|
|
"seed": args.seed,
|
|
"compile": args.compile,
|
|
"moe_mode": args.moe_mode,
|
|
"total_steps": total_steps,
|
|
"vocab_size": vocab_size,
|
|
"pinyin_vocab_size": pinyin_vocab_size,
|
|
"dim": dim,
|
|
"num_slots": num_slots,
|
|
"n_layers": n_layers,
|
|
"n_heads": n_heads,
|
|
"num_experts": num_experts,
|
|
"max_seq_len": max_seq_len,
|
|
"is_train_preprocessed": is_train_preprocessed,
|
|
"is_eval_preprocessed": is_eval_preprocessed,
|
|
}
|
|
config_file = output_path / "training_config.json"
|
|
with open(config_file, "w", encoding="utf-8") as f:
|
|
json.dump(config, f, indent=2, ensure_ascii=False)
|
|
logger.info(f"Configuration saved to {config_file}")
|
|
|
|
# ================================================================
|
|
# 创建 Trainer 并开始训练
|
|
# ================================================================
|
|
console.print("[bold cyan]正在创建训练器...[/bold cyan]")
|
|
trainer = Trainer(
|
|
model=model,
|
|
train_dataloader=train_dataloader,
|
|
eval_dataloader=eval_dataloader,
|
|
total_steps=total_steps,
|
|
output_dir=args.output_dir,
|
|
num_epochs=args.num_epochs,
|
|
learning_rate=args.learning_rate,
|
|
min_learning_rate=args.min_learning_rate,
|
|
weight_decay=args.weight_decay,
|
|
warmup_ratio=args.warmup_ratio,
|
|
label_smoothing=args.label_smoothing,
|
|
grad_accum_steps=args.grad_accum_steps,
|
|
clip_grad_norm=args.clip_grad_norm,
|
|
eval_frequency=args.eval_frequency,
|
|
save_frequency=args.save_frequency,
|
|
mixed_precision=args.mixed_precision,
|
|
device=device,
|
|
use_tensorboard=args.tensorboard,
|
|
status_file="training_status.json",
|
|
)
|
|
console.print("[green]✓ 训练器创建完成[/green]")
|
|
|
|
console.print("\n[bold cyan]开始训练...[/bold cyan]")
|
|
console.print(f"开始时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
|
|
|
try:
|
|
trainer.train(
|
|
resume_from=None,
|
|
reset_training_state=False,
|
|
auto_resume=False, # 迁移学习从头开始，不自动恢复
|
|
)
|
|
except KeyboardInterrupt:
|
|
console.print("[bold green]训练被终止[/bold green]")
|
|
trainer.save_checkpoint("interrupted_model.pt")
|
|
|
|
console.print("[bold green]✓ 训练完成![/bold green]")
|
|
console.print(f"结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
|
console.print(f"模型和日志保存在: {args.output_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|