916 lines
32 KiB
Markdown
916 lines
32 KiB
Markdown
# 输入法预测模型架构设计 (Input Method Prediction Model)
|
||
|
||
## 1. 概述
|
||
本项目旨在构建一个轻量级、高精度的中文输入法预测模型。核心设计理念是通过**结构化槽位记忆**与**交叉注意力机制**，将当前语境（光标前后文本+拼音）与历史输入习惯深度融合。为了在有限的计算资源下保持高表达能力，模型引入了**混合专家网络 (MoE)** 模块。
|
||
|
||
## 2. 核心架构流程
|
||
数据流遵循以下路径：
|
||
`输入编码` → `Transformer 上下文编码` → `槽位记忆嵌入` → `交叉注意力融合` → `门控+专家混合 (MoE)` → `分类预测` → `束搜索解码`
|
||
|
||
### 2.1 输入层设计
|
||
模型接收三类输入，分别处理以保持语义清晰：
|
||
1. **当前文本上下文**：包含光标前文本（Prefix）和光标后文本（Suffix）。
|
||
2. **拼音序列**：与当前文本对应的拼音信息，作为增强特征融入文本编码。
|
||
3. **历史槽位序列**：最近 N 个历史输入词汇，作为结构化记忆输入。
|
||
|
||
### 2.2 模块详解
|
||
|
||
#### A. Transformer 编码器 (Context Encoder)
|
||
负责提取当前语境的深层语义表示。
|
||
* **输入处理**：将 Prefix、Suffix 及拼音通过 Embedding 层映射。拼音采用**特征叠加**或**独立 Token** 方式融入，避免双流架构的复杂性。
|
||
* **骨干网络**：使用标准的 Transformer Encoder。
|
||
* **隐藏层维度**：512 [1]
|
||
* **Transformer 层数**：4 层（轻量级设计，从头训练） [1]
|
||
* **注意力头数**：4 头 [1]
|
||
* **输出**：上下文表示 $H$，形状为 `[batch, L, 512]` [1]。
|
||
|
||
#### B. 槽位记忆模块 (Slot Memory)
|
||
负责将非结构化的历史输入转化为结构化的记忆向量。
|
||
* **嵌入方式**：历史词汇通过独立的 `Slot Embedding` 查找表映射。
|
||
* **位置编码**：添加可学习的 `Positional Embedding` 以保留历史输入的时间顺序信息。
|
||
* **输出**：槽位序列 $S$，形状为 `[batch, Num_Slots, 512]`。
|
||
|
||
#### C. 交叉注意力融合 (Cross-Attention Fusion)
|
||
这是模型的核心创新点，用于动态关联"历史记忆"与"当前语境"。
|
||
* **Query (Q)**：当前步的槽位序列 $S$（经过位置编码后）。
|
||
* **Key/Value (K/V)**：Transformer 编码器输出的上下文表示 $H$ [1]。
|
||
* **机制**：让历史槽位主动关注当前文本语境，捕捉如"在'班级第一名'语境下，'王次香'比'王慈祥'更相关"的逻辑。
|
||
* **输出**：融合后的特征序列，形状为 `[batch, Num_Slots, 512]`。
|
||
|
||
#### D. 门控与专家混合 (Gating + MoE)
|
||
实际测试表明，移除 MoE 会导致模型性能显著下降，因此该模块对于捕捉复杂分布至关重要。
|
||
* **专家数量**：20 个专家 [1]。
|
||
* **门控机制**：根据输入特征动态选择激活部分专家，实现稀疏激活，在增加模型容量的同时控制计算成本。
|
||
* **输出**：经过专家网络增强后的特征向量。
|
||
|
||
#### E. 分类头与解码
|
||
* **分类预测**：MoE 输出的特征向量通过全连接层映射到词表空间，输出下一个字/词的概率分布。
|
||
* **解码策略**：推理阶段使用**束搜索 (Beam Search)**，束宽设为 5 [1]。
|
||
|
||
## 3. 关键超参数配置
|
||
|
||
为确保模型性能与效率的平衡，建议采用以下超参数 [1]：
|
||
|
||
| 参数项 | 推荐值 | 说明 |
|
||
| :--- | :--- | :--- |
|
||
| **序列长度 (L)** | 128 | 上下文窗口大小 [1] |
|
||
| **隐藏层维度** | 512 | Embedding 及 Transformer 内部维度 [1] |
|
||
| **Transformer 层数** | 4 | 轻量级骨干，降低延迟 [1] |
|
||
| **注意力头数** | 4 | 适配 512 维度的高效配置 [1] |
|
||
| **专家数量** | 20 | MoE 层中的专家总数，对性能至关重要 [1] |
|
||
| **束宽 (Beam Width)** | 5 | 推理时平衡速度与准确率 [1] |
|
||
| **学习率** | 1e-4 ~ 5e-4 | 建议配合 Warmup 策略 [1] |
|
||
|
||
## 4. 训练策略
|
||
|
||
本模型采用标准的**序列到序列（Seq2Seq）监督学习**范式，直接对目标槽位序列进行逐步预测。
|
||
|
||
### 4.1 数据构造与标签
|
||
* **输入三元组**：训练数据由 `(上下文, 拼音, 目标槽位序列)` 构成 [1]。
|
||
* **上下文**：光标前后的文本片段。
|
||
* **拼音**：当前待输入字的拼音序列。
|
||
* **目标槽位序列**：真实用户输入的文字 ID 序列，作为模型的监督信号 [1]。
|
||
* **标签处理**：在每一个槽位步（Step），模型需要预测该步对应的真实文字 ID [1]。
|
||
|
||
### 4.2 损失函数与优化
|
||
* **损失函数**：使用 **CrossEntropyLoss** 计算每一步预测结果与真实标签之间的差异 [1]。
|
||
* **掩码机制**：仅计算非填充位置（Non-padding positions）的损失，忽略无效的时间步 [1]。
|
||
* **优化器**：采用 **AdamW** 进行参数更新 [1]。
|
||
|
||
### 4.3 训练流程细节
|
||
1. **前向传播**：
|
||
* 模型接收上下文和拼音，通过 Transformer 编码得到语境表示。
|
||
* 结合历史槽位记忆，通过交叉注意力和 MoE 模块融合特征。
|
||
* 分类头输出当前步所有候选字的概率分布。
|
||
2. **Teacher Forcing**：
|
||
* 在训练过程中，**强制使用真实的上一槽位输出**作为下一步的输入条件。这意味着模型在训练时始终基于"正确的历史"进行预测，从而快速收敛。
|
||
3. **反向传播**：
|
||
* 根据 CrossEntropyLoss [1] 计算梯度，并通过 AdamW [1] 更新模型权重。
|
||
|
||
### 4.4 推理与训练的差异
|
||
* **训练时**：使用 Ground Truth（真实标签）作为槽位输入，确保模型学习到最优的条件概率分布。
|
||
* **推理时**：由于无法获取真实标签，模型采用**束搜索（Beam Search）** [1]。
|
||
* **束宽**：默认为 5 [1]。
|
||
* **候选维护**：每个候选路径独立维护其历史槽位序列及累计概率 [1]。
|
||
* **终止条件**：当所有槽位填满（如 8×3=24 步）或所有候选分支的最高概率词均为终止符时退出 [1]。
|
||
|
||
## 5. Jupyter Lab 训练示例
|
||
|
||
以下是在 Jupyter Lab 环境中使用 `trainer.Trainer` 类训练输入法模型的完整示例：
|
||
|
||
```python
|
||
# %% [markdown]
|
||
# # 输入法模型训练示例
|
||
# 本笔记本展示如何使用 trainer.Trainer 类训练输入法模型
|
||
|
||
# %% [code]
|
||
# 1. 导入必要的库
|
||
import sys
|
||
import os
|
||
from pathlib import Path
|
||
from datetime import datetime
|
||
|
||
import torch
|
||
from torch.utils.data import DataLoader
|
||
|
||
# 添加项目路径（适应不同的Jupyter Lab运行位置）
|
||
project_root = Path.cwd()
|
||
# 检查当前目录是否包含src目录，如果不包含则使用父目录
|
||
if not (project_root / "src").exists():
|
||
project_root = project_root.parent
|
||
sys.path.insert(0, str(project_root)) # 优先搜索项目目录
|
||
|
||
# 导入项目模块
|
||
from src.model.model import InputMethodEngine
|
||
from src.model.dataset import PinyinInputDataset
|
||
from src.model.trainer import Trainer, worker_init_fn, collate_fn
|
||
|
||
# %% [code]
|
||
# 2. 配置训练参数
|
||
config = {
|
||
# 数据参数
|
||
"train_data_path": "/path/to/your/train/dataset", # 替换为训练数据集路径
|
||
"eval_data_path": "/path/to/your/eval/dataset", # 替换为评估数据集路径
|
||
"output_dir": "./training_output",
|
||
|
||
# 模型参数
|
||
"vocab_size": 10019,
|
||
"pinyin_vocab_size": 30,
|
||
"dim": 512,
|
||
"num_slots": 8,
|
||
"n_layers": 4,
|
||
"n_heads": 4,
|
||
"num_experts": 20,
|
||
"max_seq_len": 128,
|
||
|
||
# 训练参数
|
||
"batch_size": 64, # 根据GPU内存调整
|
||
"num_epochs": 10,
|
||
"learning_rate": 3e-4,
|
||
"min_learning_rate": 1e-9,
|
||
"weight_decay": 0.1,
|
||
"warmup_ratio": 0.1,
|
||
"label_smoothing": 0.15,
|
||
"grad_accum_steps": 2, # 梯度累积，模拟更大batch size
|
||
"clip_grad_norm": 1.0,
|
||
"eval_frequency": 500, # 每500步评估一次
|
||
"save_frequency": 2000, # 每2000步保存检查点
|
||
|
||
# 高级选项
|
||
"mixed_precision": True,
|
||
"use_tensorboard": True,
|
||
"seed": 42,
|
||
"max_iter_length": 1024 * 1024 * 128, # 最大迭代长度
|
||
}
|
||
|
||
# %% [code]
|
||
# 3. 设置随机种子和设备
|
||
torch.manual_seed(config["seed"])
|
||
if torch.cuda.is_available():
|
||
torch.cuda.manual_seed_all(config["seed"])
|
||
device = torch.device("cuda")
|
||
print(f"✅ 使用 GPU: {torch.cuda.get_device_name(0)}")
|
||
else:
|
||
device = torch.device("cpu")
|
||
print("⚠️ 使用 CPU 进行训练（建议使用 GPU 以获得更好性能）")
|
||
|
||
# %% [code]
|
||
# 4. 创建数据集和数据加载器
|
||
print("📊 创建数据集和数据加载器...")
|
||
|
||
# 训练数据集
|
||
train_dataset = PinyinInputDataset(
|
||
data_path=config["train_data_path"],
|
||
max_workers=-1, # 自动选择worker数量
|
||
max_iter_length=config["max_iter_length"],
|
||
max_seq_length=config["max_seq_len"],
|
||
text_field="text",
|
||
py_style_weight=(9, 2, 1),
|
||
shuffle_buffer_size=5000,
|
||
length_weights={1: 10, 2: 50, 3: 50, 4: 40, 5: 15, 6: 10, 7: 5, 8: 2},
|
||
)
|
||
|
||
# 训练数据加载器
|
||
train_dataloader = DataLoader(
|
||
train_dataset,
|
||
batch_size=config["batch_size"],
|
||
num_workers=min(max(1, (os.cpu_count() or 1) - 1), 8), # 合理数量的worker
|
||
pin_memory=torch.cuda.is_available(),
|
||
worker_init_fn=worker_init_fn,
|
||
collate_fn=collate_fn,
|
||
prefetch_factor=32,
|
||
persistent_workers=True,
|
||
)
|
||
|
||
# 评估数据集
|
||
eval_dataset = PinyinInputDataset(
|
||
data_path=config["eval_data_path"],
|
||
max_workers=-1,
|
||
max_iter_length=1024, # 评估集较小
|
||
max_seq_length=config["max_seq_len"],
|
||
text_field="text",
|
||
py_style_weight=(9, 2, 1),
|
||
shuffle_buffer_size=1000,
|
||
length_weights={1: 10, 2: 50, 3: 50, 4: 40, 5: 15, 6: 10, 7: 5, 8: 2},
|
||
)
|
||
|
||
eval_dataloader = DataLoader(
|
||
eval_dataset,
|
||
batch_size=config["batch_size"],
|
||
num_workers=1,
|
||
pin_memory=torch.cuda.is_available(),
|
||
worker_init_fn=worker_init_fn,
|
||
collate_fn=collate_fn,
|
||
prefetch_factor=32,
|
||
persistent_workers=True,
|
||
)
|
||
|
||
print(f"✅ 数据加载器创建完成")
|
||
print(f" 训练批次大小: {config['batch_size']}")
|
||
print(f" 预估训练步数: {config['max_iter_length'] // config['batch_size']}")
|
||
|
||
# %% [code]
|
||
# 5. 创建模型
|
||
print("🧠 创建输入法模型...")
|
||
|
||
model = InputMethodEngine(
|
||
vocab_size=config["vocab_size"],
|
||
pinyin_vocab_size=config["pinyin_vocab_size"],
|
||
dim=config["dim"],
|
||
num_slots=config["num_slots"],
|
||
n_layers=config["n_layers"],
|
||
n_heads=config["n_heads"],
|
||
num_experts=config["num_experts"],
|
||
max_seq_len=config["max_seq_len"],
|
||
)
|
||
|
||
# 将模型移动到设备
|
||
model.to(device)
|
||
|
||
# 计算参数量
|
||
total_params = sum(p.numel() for p in model.parameters())
|
||
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
|
||
print(f"✅ 模型创建完成")
|
||
print(f" 总参数量: {total_params:,}")
|
||
print(f" 可训练参数量: {trainable_params:,}")
|
||
print(f" 模型架构: {config['n_layers']}层Transformer, {config['dim']}维度, {config['num_experts']}个MoE专家")
|
||
|
||
# %% [code]
|
||
# 6. 创建训练器
|
||
print("⚙️ 创建训练器...")
|
||
|
||
# 计算总训练步数
|
||
total_steps = int(config["max_iter_length"] / config["batch_size"])
|
||
|
||
trainer = Trainer(
|
||
model=model,
|
||
train_dataloader=train_dataloader,
|
||
eval_dataloader=eval_dataloader,
|
||
total_steps=total_steps,
|
||
output_dir=config["output_dir"],
|
||
num_epochs=config["num_epochs"],
|
||
learning_rate=config["learning_rate"],
|
||
min_learning_rate=config["min_learning_rate"],
|
||
weight_decay=config["weight_decay"],
|
||
warmup_ratio=config["warmup_ratio"],
|
||
label_smoothing=config["label_smoothing"],
|
||
grad_accum_steps=config["grad_accum_steps"],
|
||
clip_grad_norm=config["clip_grad_norm"],
|
||
eval_frequency=config["eval_frequency"],
|
||
save_frequency=config["save_frequency"],
|
||
mixed_precision=config["mixed_precision"],
|
||
use_tensorboard=config["use_tensorboard"],
|
||
)
|
||
|
||
print(f"✅ 训练器创建完成")
|
||
print(f" 总训练步数: {total_steps:,}")
|
||
print(f" 学习率: {config['learning_rate']:.2e} -> {config['min_learning_rate']:.2e}")
|
||
print(f" 输出目录: {config['output_dir']}")
|
||
|
||
# %% [code]
|
||
# 7. 开始训练
|
||
print("🚀 开始训练...")
|
||
print(f"开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
|
||
try:
|
||
# 开始训练（可以从检查点恢复训练）
|
||
trainer.train(resume_from=None) # 设置检查点路径以恢复训练
|
||
|
||
print("✅ 训练完成!")
|
||
print(f"结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
print(f"模型和日志保存在: {config['output_dir']}")
|
||
|
||
except KeyboardInterrupt:
|
||
print("⏹️ 训练被用户中断")
|
||
print("💾 保存当前检查点...")
|
||
trainer.save_checkpoint("interrupted")
|
||
print(f"检查点已保存到: {config['output_dir']}/checkpoint_interrupted.pt")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 训练过程中出现错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
# %% [code]
|
||
# 8. 监控训练进度（如果使用TensorBoard）
|
||
if config["use_tensorboard"]:
|
||
print("📈 TensorBoard日志已记录在:")
|
||
print(f" {config['output_dir']}/tensorboard")
|
||
print("\n启动TensorBoard查看训练进度:")
|
||
print(" tensorboard --logdir ./training_output/tensorboard")
|
||
print("然后在浏览器中打开: http://localhost:6006")
|
||
|
||
# %% [code]
|
||
# 9. 加载训练好的模型进行推理（示例）
|
||
def load_trained_model(checkpoint_path):
|
||
"""加载训练好的模型进行检查点"""
|
||
print(f"📥 加载检查点: {checkpoint_path}")
|
||
|
||
# 创建与训练时相同配置的模型
|
||
loaded_model = InputMethodEngine(
|
||
vocab_size=config["vocab_size"],
|
||
pinyin_vocab_size=config["pinyin_vocab_size"],
|
||
dim=config["dim"],
|
||
num_slots=config["num_slots"],
|
||
n_layers=config["n_layers"],
|
||
n_heads=config["n_heads"],
|
||
num_experts=config["num_experts"],
|
||
max_seq_len=config["max_seq_len"],
|
||
)
|
||
|
||
# 加载检查点
|
||
checkpoint = torch.load(checkpoint_path, map_location=device)
|
||
loaded_model.load_state_dict(checkpoint["model_state_dict"])
|
||
loaded_model.to(device)
|
||
loaded_model.eval()
|
||
|
||
print(f"✅ 模型加载完成，训练步数: {checkpoint.get('global_step', 'N/A')}")
|
||
print(f" 训练损失: {checkpoint.get('train_loss', 'N/A'):.4f}")
|
||
|
||
return loaded_model
|
||
|
||
# 使用示例（取消注释以使用）
|
||
# trained_model = load_trained_model("./training_output/checkpoint_final.pt")
|
||
```
|
||
|
||
### 关键说明
|
||
|
||
1. **环境要求**：
|
||
- Python 3.12+
|
||
- PyTorch 2.10+
|
||
- 建议使用GPU进行训练
|
||
- 安装项目依赖：`pip install -e .`
|
||
|
||
2. **数据集格式**：
|
||
- 使用Hugging Face `datasets`格式
|
||
- 必须包含`text`字段
|
||
- 支持流式读取（streaming=True）
|
||
|
||
3. **训练监控**：
|
||
- 控制台输出训练进度和指标
|
||
- TensorBoard记录损失、准确率、学习率等
|
||
- 定期保存模型检查点
|
||
|
||
4. **可调整参数**：
|
||
- `batch_size`: 根据GPU内存调整
|
||
- `learning_rate`: 建议在1e-4到5e-4之间
|
||
- `grad_accum_steps`: 模拟更大batch size
|
||
- `num_epochs`: 根据数据集大小调整
|
||
|
||
5. **故障排除**：
|
||
- GPU内存不足：减小`batch_size`或增加`grad_accum_steps`
|
||
- 训练不稳定：降低`learning_rate`或增加`warmup_ratio`
|
||
- 过拟合：增加`label_smoothing`或使用更大数据集
|
||
|
||
## 6. 使用指南
|
||
|
||
本项目的训练功能通过命令行工具 `train-model` 提供，支持训练、评估和导出模型。
|
||
|
||
### 6.1 安装与准备
|
||
|
||
#### 使用 uv（推荐）
|
||
本项目使用 [`uv`](https://github.com/astral-sh/uv) 作为Python包管理器，它比传统的 pip 更快且更可靠。
|
||
|
||
1. **安装 uv**（如果尚未安装）：
|
||
```bash
|
||
# Linux/macOS
|
||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||
|
||
# 或使用 pipx
|
||
pipx install uv
|
||
|
||
# Windows (PowerShell)
|
||
powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
|
||
```
|
||
|
||
2. **安装项目依赖**：
|
||
```bash
|
||
uv pip install -e .
|
||
```
|
||
|
||
#### 使用传统 pip
|
||
如果不使用 uv，也可以用标准的 pip 安装：
|
||
|
||
```bash
|
||
# 创建并激活虚拟环境（推荐）
|
||
python -m venv .venv
|
||
source .venv/bin/activate # Linux/macOS
|
||
# .venv\Scripts\activate # Windows
|
||
|
||
# 安装依赖
|
||
pip install -e .
|
||
```
|
||
|
||
#### 验证安装
|
||
安装完成后，可通过以下命令验证：
|
||
```bash
|
||
train-model --help
|
||
```
|
||
|
||
### 6.2 数据格式
|
||
|
||
训练数据应为Hugging Face数据集格式，支持本地文件或远程数据集仓库。数据集需包含 `text` 字段，并支持流式读取（streaming=True）。
|
||
|
||
#### 本地数据集示例
|
||
```python
|
||
# dataset.py
|
||
from datasets import Dataset
|
||
|
||
data = {
|
||
"text": ["这是第一个样本文本。", "这是第二个样本，用于训练输入法模型。"]
|
||
}
|
||
dataset = Dataset.from_dict(data)
|
||
dataset.save_to_disk("./local_dataset")
|
||
```
|
||
|
||
#### 远程数据集示例
|
||
支持Hugging Face Hub或ModelScope上的数据集：
|
||
- `huggingface.co/datasets/username/dataset_name`
|
||
- `modelscope.cn/datasets/username/dataset_name`
|
||
|
||
#### 数据格式要求
|
||
- **必需字段**: `text`（字符串类型，包含中文文本）
|
||
- **流式读取**: 数据集必须支持 `streaming=True` 参数
|
||
- **数据量**: 建议至少数百万条文本以获得良好效果
|
||
|
||
#### 数据预处理
|
||
数据集会自动进行以下处理：
|
||
1. 文本分词和编码
|
||
2. 拼音转换和编码
|
||
3. 上下文窗口滑动生成训练样本
|
||
4. 频率调整（削峰填谷）以平衡高频/低频字词
|
||
|
||
### 6.3 基本训练命令
|
||
|
||
使用 `train-model train` 命令开始训练：
|
||
|
||
```bash
|
||
train-model train \
|
||
--train-data-path "path/to/train/dataset" \
|
||
--eval-data-path "path/to/eval/dataset" \
|
||
--output-dir "./output" \
|
||
--batch-size 128 \
|
||
--num-epochs 10 \
|
||
--learning-rate 1e-5
|
||
```
|
||
|
||
#### 检查点恢复训练
|
||
|
||
要从检查点恢复训练（保持原有的训练状态）：
|
||
|
||
```bash
|
||
train-model train \
|
||
--train-data-path "path/to/train/dataset" \
|
||
--eval-data-path "path/to/eval/dataset" \
|
||
--resume-from "./output/checkpoints/latest_checkpoint.pt"
|
||
```
|
||
|
||
#### 重置训练状态
|
||
|
||
如果只想加载模型权重，从头开始训练（学习率、epoch等都重新开始）：
|
||
|
||
```bash
|
||
train-model train \
|
||
--train-data-path "path/to/train/dataset" \
|
||
--eval-data-path "path/to/eval/dataset" \
|
||
--resume-from "./output/checkpoints/best_model.pt" \
|
||
--reset-training-state
|
||
```
|
||
|
||
这个功能在以下场景非常有用：
|
||
- 想要用预训练权重初始化模型，但用新的训练计划重新训练
|
||
- 需要调整学习率策略或训练时长
|
||
- 在现有模型基础上进行迁移学习
|
||
|
||
#### 学习率建议
|
||
根据模型架构和超参数配置（4层Transformer，512维度），推荐使用以下学习率范围：
|
||
- **标准范围**: 1e-4 ~ 5e-4
|
||
- **配合Warmup策略**：在训练初期逐步提高学习率
|
||
- **余弦退火**：使用最小学习率 1e-9 进行细调
|
||
|
||
### 6.4 参数详解
|
||
|
||
#### 数据参数
|
||
- `--train-data-path`, `-t`: 训练数据集路径（必需）
|
||
- `--eval-data-path`, `-e`: 评估数据集路径（必需）
|
||
- `--output-dir`, `-o`: 输出目录（默认：`./output`）
|
||
- `--max_iter_length`: 最大迭代长度，控制每次训练迭代处理的数据量（默认：134217728）
|
||
|
||
#### 模型参数
|
||
- `--vocab-size`: 词汇表大小（默认：10019）
|
||
- `--pinyin-vocab-size`: 拼音词汇表大小（默认：30）
|
||
- `--dim`: 模型维度（默认：512）
|
||
- `--num-slots`: 历史槽位数量（默认：8）
|
||
- `--n-layers`: Transformer层数（默认：4）
|
||
- `--n-heads`: 注意力头数（默认：4）
|
||
- `--num-experts`: MoE专家数量（默认：20）
|
||
- `--max-seq-len`: 最大序列长度（默认：128）
|
||
- `--use-pinyin`: 是否使用拼音特征（默认：False）
|
||
|
||
#### 训练参数
|
||
- `--batch-size`, `-b`: 批次大小（默认：128）
|
||
- `--num-epochs`: 训练轮数（默认：10）
|
||
- `--learning-rate`, `-lr`: 学习率（默认：1e-5）
|
||
- `--min-learning-rate`: 最小学习率（默认：1e-9）
|
||
- `--weight-decay`: 权重衰减（默认：0.1）
|
||
- `--warmup-ratio`: 热身步数比例（默认：0.1）
|
||
- `--label-smoothing`: 标签平滑参数（默认：0.15）
|
||
- `--grad-accum-steps`: 梯度累积步数（默认：1）
|
||
- `--clip-grad-norm`: 梯度裁剪范数（默认：1.0）
|
||
- `--eval-frequency`: 评估频率（默认：500步）
|
||
- `--save-frequency`: 保存频率（默认：10000步）
|
||
|
||
#### 高级选项
|
||
- `--mixed-precision/--no-mixed-precision`: 是否使用混合精度训练（默认：启用）
|
||
- `--tensorboard/--no-tensorboard`: 是否使用TensorBoard（默认：启用）
|
||
- `--resume-from`: 从检查点恢复训练（可选）
|
||
- `--reset-training-state`: 重置训练状态，只加载模型权重从头开始训练（默认：False）
|
||
- `--seed`: 随机种子（默认：42）
|
||
|
||
### 6.5 监控训练进度
|
||
|
||
训练过程中会显示：
|
||
- 当前训练步数/总步数
|
||
- 损失值和准确率
|
||
- 学习率变化
|
||
- 内存使用情况
|
||
|
||
启用TensorBoard后，可使用以下命令查看可视化结果：
|
||
|
||
```bash
|
||
tensorboard --logdir ./output/tensorboard
|
||
```
|
||
|
||
### 6.6 基于JSON旁路记录法的移动端监控方案
|
||
|
||
为了提供移动端友好的训练监控体验，我们实现了基于JSON旁路记录法的监控方案。该方案在保持TensorBoard记录的同时，额外写入一份JSON状态文件，并通过Flask提供移动端友好的Web界面。
|
||
|
||
#### 方案特点
|
||
|
||
**📱 移动端体验**
|
||
- Flask模板生成响应式界面，完美适配手机屏幕
|
||
- 图表支持双指缩放和滑动操作
|
||
- 大字体显示核心指标，触控操作便捷
|
||
|
||
**🚀 低耦合架构**
|
||
- 训练和监控通过文件系统解耦
|
||
- 监控服务重启不影响训练进程
|
||
- 训练脚本只需几行代码修改即可支持
|
||
|
||
**🔒 安全稳定**
|
||
- 纯文本JSON文件，无文件锁冲突问题
|
||
- 读写速度快，稳定性高
|
||
- 不会影响原有的TensorBoard记录流程
|
||
|
||
**📊 实时监控**
|
||
- 默认每5秒自动刷新数据
|
||
- 实时显示训练进度和指标趋势
|
||
- 数据新鲜度状态指示（实时/较新/较旧/陈旧）
|
||
|
||
#### 使用方法
|
||
|
||
**启动监控服务**
|
||
```bash
|
||
# 启动监控服务（默认端口8501）
|
||
monitor-training monitor
|
||
|
||
# 指定状态文件路径和端口
|
||
monitor-training monitor --status-file ./output/training_status.json --port 8080
|
||
|
||
# 不自动打开浏览器
|
||
monitor-training monitor --no-browser
|
||
|
||
# 指定自定义Flask应用脚本
|
||
monitor-training monitor --streamlit-script ./custom_monitor.py
|
||
```
|
||
|
||
**查看训练状态**
|
||
```bash
|
||
# 查看最近10条训练记录
|
||
monitor-training view
|
||
|
||
# 查看最近50条记录（原始JSON格式）
|
||
monitor-training view --limit 50 --raw
|
||
|
||
# 查看指定状态文件
|
||
monitor-training view /path/to/status.json
|
||
```
|
||
|
||
**检查状态文件**
|
||
```bash
|
||
# 检查状态文件状态
|
||
monitor-training check
|
||
|
||
# 检查指定文件
|
||
monitor-training check ./output/training_status.json
|
||
```
|
||
|
||
**启动HTTP静态文件服务**
|
||
```bash
|
||
# 启动HTTP静态文件服务（默认端口8080）
|
||
monitor-training serve
|
||
|
||
# 指定状态文件路径和端口
|
||
monitor-training serve --status-file ./output/training_status.json --port 8080
|
||
|
||
# 禁用CORS支持（默认启用）
|
||
monitor-training serve --no-cors
|
||
|
||
# 指定主机地址
|
||
monitor-training serve --host 0.0.0.0 --port 8080
|
||
```
|
||
|
||
#### 监控界面功能
|
||
|
||
**📊 核心指标看板**
|
||
- 当前步数、轮次、训练损失、准确率
|
||
- 评估损失和准确率
|
||
- 当前学习率、最后更新时间
|
||
|
||
**📈 趋势图表**
|
||
- 损失曲线（训练损失 + 评估损失）
|
||
- 准确率曲线（训练准确率 + 评估准确率）
|
||
- 学习率变化图（对数坐标）
|
||
|
||
**📋 数据详情**
|
||
- 完整的训练记录表格
|
||
- 数据统计信息（总数据点、训练时长、总步数）
|
||
- 训练进度条
|
||
|
||
**⚙️ 配置选项**
|
||
- 状态文件路径（支持环境变量 `TRAINING_STATUS_FILE`）
|
||
- 自动刷新间隔（1-30秒可调）
|
||
- 显示数据点数量（10-1000条可调）
|
||
|
||
#### 技术实现
|
||
|
||
**训练端改造**
|
||
- 在 `Trainer.__init__` 中添加 `status_file` 参数
|
||
- 实现 `_write_training_status()` 方法，在每次评估时写入JSON文件
|
||
- 支持从现有状态文件恢复，避免数据丢失
|
||
|
||
**监控端搭建**
|
||
- 使用Flask构建移动端友好的Web界面
|
||
- 采用Plotly图表库，支持触控交互
|
||
- 自动刷新机制，实时更新训练状态
|
||
|
||
**命令行工具**
|
||
- 提供 `monitor`、`view`、`check` 三个子命令
|
||
- 自动检测Flask可用性
|
||
- 支持环境变量传递
|
||
|
||
#### 访问方式
|
||
|
||
**本地访问**
|
||
```bash
|
||
# 启动监控服务后，通过浏览器访问
|
||
http://localhost:8501
|
||
```
|
||
|
||
**局域网访问**
|
||
```bash
|
||
# 启动服务时指定主机地址
|
||
monitor-training monitor --host 0.0.0.0 --port 8080
|
||
|
||
# 手机浏览器访问（同一局域网）
|
||
http://192.168.1.100:8080
|
||
```
|
||
|
||
**公网访问**（需端口转发）
|
||
```bash
|
||
# 确保服务器防火墙开放对应端口
|
||
# 通过域名或公网IP访问
|
||
http://your-server.com:8501
|
||
```
|
||
|
||
**远程HTTP监控**
|
||
```bash
|
||
# GPU服务器启动HTTP服务
|
||
monitor-training serve --port 8080 --host 0.0.0.0
|
||
|
||
# 本地运行Flask监控，从HTTP URL读取数据
|
||
monitor-training monitor --host 127.0.0.1 --port 8501
|
||
|
||
# 在Flask界面输入远程URL：
|
||
http://<gpu服务器IP>:8080/training_status.json
|
||
```
|
||
|
||
#### 状态文件格式
|
||
|
||
状态文件 `training_status.json` 位于训练输出目录，格式如下：
|
||
```json
|
||
[
|
||
{
|
||
"step": 100,
|
||
"epoch": 1,
|
||
"timestamp": "2024-01-01T12:00:00",
|
||
"train/loss": 2.345,
|
||
"train/accuracy": 0.456,
|
||
"eval/loss": 2.123,
|
||
"eval/accuracy": 0.512,
|
||
"train/learning_rate": 0.0001
|
||
},
|
||
...
|
||
]
|
||
```
|
||
|
||
#### HTTP静态文件服务与远程监控
|
||
|
||
针对GPU服务器只支持HTTP协议（不支持WebSockets）的环境，我们提供了HTTP静态文件服务方案，实现远程训练监控。
|
||
|
||
**🔧 技术特点**
|
||
- 纯HTTP协议，无需WebSockets支持
|
||
- 原子写入机制，避免读取不完整JSON数据
|
||
- 自动重试和JSON验证，确保数据完整性
|
||
- CORS支持，方便跨域访问
|
||
- 轻量级设计，不影响训练性能
|
||
|
||
**🚀 工作原理**
|
||
1. GPU服务器：训练进程通过原子写入机制更新`training_status.json`文件
|
||
2. GPU服务器：运行`monitor-training serve`提供HTTP静态文件服务
|
||
3. 本地机器：运行`monitor-training monitor`启动Flask监控界面
|
||
4. 本地机器：在Flask界面输入HTTP URL访问远程数据
|
||
5. Flask：通过HTTP轮询获取实时训练数据并展示
|
||
|
||
**🛡️ 数据安全**
|
||
- 原子写入：先写入临时文件，然后原子重命名，避免读取中断
|
||
- JSON验证：HTTP服务端验证JSON格式后才返回数据
|
||
- 临时文件处理：智能识别和读取`.tmp`临时文件
|
||
- 重试机制：JSON解析失败时自动重试读取
|
||
|
||
**🌐 网络要求**
|
||
- GPU服务器：需要开放HTTP端口（默认8080）
|
||
- 本地机器：需要能访问GPU服务器的HTTP端口
|
||
- 网络协议：纯HTTP，兼容防火墙和代理
|
||
|
||
#### 注意事项
|
||
1. 首次监控时如果状态文件不存在，会自动创建空文件
|
||
2. 需要安装 `plotly` 依赖用于图表绘制：`pip install plotly>=5.0.0`
|
||
3. 从检查点恢复训练时会自动加载已有的状态数据
|
||
4. 建议将监控服务与训练服务部署在同一服务器，避免网络延迟
|
||
5. HTTP服务支持原子写入，避免训练进程写入时读取不完整JSON
|
||
6. 远程监控需要确保GPU服务器防火墙开放对应HTTP端口
|
||
7. 建议使用`--host 0.0.0.0`参数使HTTP服务可被远程访问
|
||
|
||
|
||
### 6.7 评估模型（开发中）
|
||
|
||
当前评估功能尚在开发中：
|
||
|
||
```bash
|
||
train-model evaluate \
|
||
--checkpoint "./output/checkpoint_final.pt" \
|
||
--data-path "path/to/eval/dataset" \
|
||
--batch-size 32
|
||
```
|
||
|
||
命令将显示"评估功能待实现"的提示信息。该功能计划用于：
|
||
- 加载训练好的模型检查点
|
||
- 在评估数据集上计算准确率、困惑度等指标
|
||
- 生成详细的性能报告
|
||
|
||
### 6.8 模型扩容两阶段训练
|
||
|
||
当需要增加模型容量（如增加专家数量、修改层结构等）时，可以使用 `expand-and-train` 命令进行两阶段训练：先冻结匹配层训练新增参数，然后全量微调。
|
||
|
||
#### 训练策略
|
||
|
||
1. **冻结阶段**：只训练形状不匹配的新增参数（如新增的专家、扩容的层等）
|
||
2. **全量微调阶段**：当验证损失连续 `--frozen-patience` 次不下降时，自动解冻所有层进行全量训练
|
||
|
||
#### 基础用法
|
||
|
||
```bash
|
||
train-model expand-and-train \
|
||
--train-data-path "path/to/train/dataset" \
|
||
--eval-data-path "path/to/eval/dataset" \
|
||
--base-model-path "./pretrained/model.pt" \
|
||
--new-model-spec "model:InputMethodEngine" \
|
||
--num-experts 40 \
|
||
--frozen-lr 2e-3 \
|
||
--full-lr 5e-5 \
|
||
--frozen-patience 8
|
||
```
|
||
|
||
#### 完整参数示例
|
||
|
||
```bash
|
||
train-model expand-and-train \
|
||
--train-data-path "path/to/train/dataset" \
|
||
--eval-data-path "path/to/eval/dataset" \
|
||
--output-dir "./expansion_output" \
|
||
--base-model-path "./pretrained/model.pt" \
|
||
--new-model-spec "custom_model:ExpandedModel" \
|
||
--vocab-size 10019 \
|
||
--dim 512 \
|
||
--num-experts 40 \
|
||
--frozen-patience 10 \
|
||
--frozen-lr 1e-3 \
|
||
--full-lr 1e-4 \
|
||
--frozen-scheduler cosine \
|
||
--full-scheduler cosine \
|
||
--batch-size 128 \
|
||
--num-epochs 20 \
|
||
--compile
|
||
```
|
||
|
||
#### 参数详解
|
||
|
||
**模型扩容参数**
|
||
- `--base-model-path`: 预训练基础模型检查点路径（必需）
|
||
- `--new-model-spec`: 新模型规格，格式：`模块名:类名`，如 `model:InputMethodEngine`（必需）
|
||
- 支持任意路径的模块导入，模块文件需包含自定义的模型类
|
||
- 自定义模型类必须是 `InputMethodEngine` 的子类
|
||
- 示例：`my_model:MyExpandedModel` 对应 `my_model.py` 中的 `MyExpandedModel` 类
|
||
|
||
**两阶段训练参数**
|
||
- `--frozen-patience`: 冻结阶段验证损失连续不下降的评估次数，触发切换到全量微调（默认：10）
|
||
- `--frozen-lr`: 冻结阶段学习率（默认：1e-3）
|
||
- `--full-lr`: 全量微调阶段学习率（默认：1e-4）
|
||
- `--frozen-scheduler`: 冻结阶段学习率调度器，可选 `cosine` 或 `plateau`（默认：`cosine`）
|
||
- `--full-scheduler`: 全量微调阶段学习率调度器，可选 `cosine` 或 `plateau`（默认：`cosine`）
|
||
|
||
**其他参数**
|
||
- 支持所有 `train` 子命令的通用参数（数据参数、模型参数、训练参数等）
|
||
- 继承现有的训练基础设施：混合精度训练、TensorBoard日志、checkpoint保存等
|
||
|
||
#### 使用场景
|
||
|
||
1. **增加专家数量**（20→40）
|
||
- 冻结效果：~70% 参数可冻结（已有专家权重、注意力层等）
|
||
- 新增参数：新专家网络、gate层
|
||
|
||
2. **增加top_k值**（2→3）
|
||
- 冻结效果：100% 参数可冻结（仅逻辑变化）
|
||
- 新增参数：无
|
||
|
||
3. **修改专家内部结构**（如增加resblocks）
|
||
- 冻结效果：~50% 参数可冻结（linear_in/output可冻结）
|
||
- 新增参数：新增的resblocks层
|
||
|
||
4. **增加Transformer层数**（4→5）
|
||
- 冻结效果：~80% 参数可冻结（前4层可冻结）
|
||
- 新增参数：新增的第5层
|
||
|
||
#### 自定义模型类示例
|
||
|
||
```python
|
||
# my_model.py
|
||
from model.model import InputMethodEngine
|
||
|
||
class MyExpandedModel(InputMethodEngine):
|
||
def __init__(self, num_experts=40, **kwargs):
|
||
# 调用父类构造函数，覆盖num_experts参数
|
||
super().__init__(num_experts=num_experts, **kwargs)
|
||
# 可以在这里添加额外的层或修改现有层
|
||
|
||
# 使用命令
|
||
# train-model expand-and-train --new-model-spec "my_model:MyExpandedModel" ...
|
||
```
|
||
|
||
#### 注意事项
|
||
|
||
1. **模型类要求**：自定义模型类必须是 `InputMethodEngine` 的子类
|
||
2. **冻结条件**：只有权重形状完全匹配的层才会被冻结
|
||
3. **性能保持**：MoE层保持"计算所有专家+Top-K选择"方案，确保 `torch.compile` 下的最佳性能
|
||
4. **阶段切换**：基于评估频率而非epoch，建议适当调高 `--eval-frequency`
|
||
5. **模块导入**：支持任意路径的模块，通过Python标准导入机制加载
|
||
|
||
### 6.9 导出模型（开发中）
|
||
|
||
当前导出功能尚在开发中：
|
||
|
||
```bash
|
||
train-model export \
|
||
--checkpoint "./output/checkpoint_final.pt" \
|
||
--output "./exported_model.onnx"
|
||
```
|
||
|
||
命令将显示"导出功能待实现"的提示信息。该功能计划用于：
|
||
- 将PyTorch模型转换为ONNX格式
|
||
- 支持在不同推理引擎上部署
|
||
- 提供优化后的推理模型
|
||
|
||
## 7. 总结
|
||
本方案通过**单流 Transformer 编码**结合**结构化槽位交叉注意力**，并引入**20个专家的 MoE 模块** [1]，在保证模型轻量（4层 Transformer）的同时，有效利用了历史输入习惯并提升了模型表达上限。相比暴力拼接或双流架构，该设计在工程实现上更优雅，在推理效率上更高效，是轻量级输入法模型的局部最优解。 |