SpeechBrain 语言模型训练模板实战：基于 RNNLM 与 HuggingFace Datasets 的文本建模指南

导读

本文以 SpeechBrain 仓库中的 templates/speech_recognition/LM 模板为核心，系统讲解如何用 RNN（LSTM）语言模型在文本语料上训练 LM，并说明它与 ASR 流水线（Tokenizer → LM → ASR）的衔接关系。读完本文，你将掌握 python train.py RNNLM.yaml 的完整运行流程、RNNLM.yaml 中每个核心超参数的语义、训练脚本 train.py 中从数据加载到 Brain 训练循环的完整调用链，以及如何把模板中的 CustomModel 替换为 SpeechBrain 内置的 RNNLM 或 TransformerLM。

模板定位：ASR 流水线中的语言模型环节

在 SpeechBrain 的端到端语音识别模板（templates/speech_recognition/README.md）中，系统训练分三步走：

1. 训练 Tokenizer（templates/speech_recognition/Tokenizer）：基于训练转录文本，用 SentencePiece 学出介于字符与整词之间的子词单元；
2. 训练语言模型（即本模板 templates/speech_recognition/LM）：在同样的 token 序列上学习词的先验概率分布；
3. 训练语音识别器（templates/speech_recognition/ASR）：用 CRDNN + 注意力解码器做端到端识别，最终在解码阶段用 beam search 耦合本模板训练出的 RNN LM 进行重打分。

因此，本模板的定位非常明确：它是 ASR 解码阶段的语言先验来源，而不是独立的识别系统。模板 README 明确提示，运行本模板前必须先完成 Tokenizer 训练（见 templates/speech_recognition/Tokenizer/README.md），因为 LM 训练直接依赖 Tokenizer 产出的 1000_unigram.model 模型文件。

为什么需要一个独立的语言模型

语言模型解决的是"给定前文，下一个词是什么"的概率建模问题。在 ASR 系统中，声学模型负责"音频 → 音素/子词"，而 LM 负责"子词序列 → 词序列"的语言合理性约束，两者结合能显著降低同音词、语法错误等识别错误。

模板 README 特别提醒了 LM 训练的计算成本：在大规模文本语料上，即使用现代 GPU，训练一个 LM 也可能耗时数周甚至数月。因此在真实场景中，常见的做法是复用已有的预训练 LM 或在其上做微调（fine-tune），而不是从零训练。本模板为了演示流程，刻意只使用 mini-librispeech 的训练转录文本（templates/speech_recognition/LM/data/train.txt，共 1217 行），让用户能快速跑通完整流程。

运行前准备：安装依赖

本模板有一个额外依赖：HuggingFace 的 datasets 库。LM 的数据读取与加载完全由它托管（详见 extra_requirements.txt）。安装命令：

pip install datasets

此外，因为训练脚本中通过 sentencepiece.SentencePieceProcessor 加载 tokenizer 模型，你还需要安装 sentencepiece（Tokenizer 训练阶段已隐含此依赖）。

如何运行

在 templates/speech_recognition/LM 目录下执行：

python train.py RNNLM.yaml

运行时，脚本会依次完成：解析命令行参数 → 初始化 DDP（多卡训练时启用）→ 加载 YAML 超参数 → 创建实验目录 → 用 Pretrainer 下载/加载 tokenizer → 构建数据集 → 构造 Brain 对象 → fit() 训练 → evaluate() 在测试集上评估。

训练脚本源码解析

train.py 是模板的核心脚本，其结构可拆成三部分：

1. 自定义 Brain 子类 LM

LM(sb.core.Brain) 继承了 SpeechBrain 的 Brain 基类 提供的训练循环框架，只需覆写四个钩子方法：

• compute_forward：从 batch 中取出 tokens_bos（带句首标记的 token 序列），喂给 self.hparams.model，返回下一词的预测分布。因为 LSTM 是自回归的，输入是"前文 + BOS"，预测目标自然对齐到"下一个 token"。
• compute_objectives：取出 tokens_eos（带句尾标记的序列）作为监督目标，调用 self.hparams.compute_cost（配置为 speechbrain.nnet.losses.nll_loss，即负对数似然损失）计算损失并返回。
• on_fit_batch_end：在每个优化步结束后，若调度器是 NoamScheduler 或 CyclicCosineScheduler，则按步更新学习率。
• on_stage_end：阶段结束时记录统计；在验证阶段调用 lr_annealing(stage_loss)（NewBob 调度）更新学习率、写日志，并通过 checkpointer.save_and_keep_only(meta=stats, min_keys=["loss"]) 保存验证损失最小的 checkpoint。

2. 数据管线 dataio_prepare

这是模板与 HuggingFace datasets 集成的关键函数：

datasets = load_dataset(
    "text",
    data_files={
        "train": hparams["lm_train_data"],
        "valid": hparams["lm_valid_data"],
        "test": hparams["lm_test_data"],
    },
)

它用 load_dataset("text", ...) 把 YAML 中指定的三个纯文本文件加载为 HF 数据集，再通过 sb.dataio.dataset.DynamicItemDataset.from_arrow_dataset(...) 转成 SpeechBrain 的 DynamicItemDataset（speechbrain/dataio/dataset.py）。

随后定义 text_pipeline 文本处理管线，用 @takes("text") / @provides("text", "tokens_bos", "tokens_eos") 声明输入输出（管线装饰器实现在 speechbrain/utils/data_pipeline.py）：

• 原始文本原样透传；
• tokens_bos = [bos_index] + tokenizer.encode_as_ids(text)：在序列开头拼上 BOS 索引，作为模型输入；
• tokens_eos = tokenizer.encode_as_ids(text) + [eos_index]：在序列结尾拼上 EOS 索引，作为损失计算的目标。

最后用 set_output_keys(datasets, ["id", "text", "tokens_bos", "tokens_eos"]) 声明每个 batch 中携带的字段，训练循环里就能直接访问 batch.tokens_bos 和 batch.tokens_eos。

注意：在模板的 RNNLM.yaml 中 bos_index 与 eos_index 都配置为 0，即复用 SentencePiece 的 BOS/EOS 特殊符号索引（<s> / </s>）。自定义数据时务必与 Tokenizer 训练时的特殊符号配置保持一致。

3. 主入口

if __name__ == "__main__": 中的流程与 SpeechBrain 其他模板一脉相承：

1. sb.parse_arguments(sys.argv[1:]) 解析命令行参数（含 --hparams 覆盖能力）；
2. sb.utils.distributed.ddp_init_group(run_opts) 初始化多卡 DDP 环境（speechbrain/utils/distributed.py）；
3. load_hyperpyyaml(fin, overrides) 加载超参数文件（speechbrain/utils/hparams.py）；
4. sb.create_experiment_directory(...) 创建实验目录并备份 YAML；
5. hparams["pretrainer"].collect_files() + load_collected() 加载 tokenizer 模型；
6. 构造 LM Brain 对象，调用 fit() 训练、evaluate() 测试。

超参数文件 RNNLM.yaml 逐项解读

RNNLM.yaml 是模板唯一给出的配置，完整覆盖了数据、训练、模型、优化器、调度器、checkpoint 六大块：

数据与输出路径

配置项	默认值	说明
seed	2602	随机种子，必须在 YAML 顶部、任何带参数对象创建之前设置
data_folder	data/	语料目录
output_folder	results/RNNLM/	实验输出目录
save_folder	results/RNNLM/save	checkpoint 保存目录
train_log	results/RNNLM/train_log.txt	训练日志文件

YAML 中还给出了一条面向 HPC 集群/大数据集 的强烈建议：把语料压缩成单个 tar/zip 后拷贝到计算节点的本地磁盘（如 SLURM 集群的 $SLURM_TMPDIR），再将 lm_train_data / lm_valid_data / lm_test_data 指向本地路径。理由是从本地磁盘读取远快于共享文件系统，可避免 I/O 成为大规模 LM 训练的瓶颈。

Tokenizer 与 Pretrainer

tokenizer_file: ../Tokenizer/save/1000_unigram.model

LM 必须复用与 ASR 训练相同的 tokenizer（模板注释明确要求）。pretrainer 对象通过 speechbrain/utils/parameter_transfer.py 中的 Pretrainer 把 1000_unigram.model 加载进 sentencepiece.SentencePieceProcessor 实例。这里路径 ../Tokenizer/save/... 表明：必须先运行 Tokenizer 模板训练出该文件。

训练参数

配置项	默认值	说明
number_of_epochs	20	训练轮数
batch_size	80	训练 batch 大小
lr	0.001	初始学习率
grad_accumulation_factor	1	梯度累积倍数，用于模拟更大的 batch
ckpt_interval_minutes	15	每 N 分钟保存一次 checkpoint（中断恢复用）

Dataloader 选项

• train_dataloader_opts：batch_size=80、shuffle=True；
• valid_dataloader_opts / test_dataloader_opts：batch_size=1（验证/测试阶段不做并行计算）。

模型参数

配置项	默认值	说明
emb_dim	256	词嵌入维度
rnn_size	512	隐层神经元数
layers	2	RNN 层数
bos_index / eos_index	0	句首/句尾标记索引

优化器与学习率调度

optimizer: !name:torch.optim.Adam
    lr: !ref <lr>
    betas: (0.9, 0.98)
    eps: 0.000000001

优化器在 Brain 类把所有参数搬到正确设备后构造，并交由 Checkpointer 管理。学习率调度采用 NewBob 算法（speechbrain/nnet/schedulers.py 中的 NewBobScheduler）：若连续两个 epoch 的验证损失改善低于阈值 improvement_threshold: 0.0025，就把学习率乘以 annealing_factor: 0.8。patient: 0 表示不等待，直接按阈值判定。

Checkpointer 与 Epoch Counter

checkpointer: !new:speechbrain.utils.checkpoints.Checkpointer
    checkpoints_dir: !ref <save_folder>
    recoverables:
        model: !ref <model>
        scheduler: !ref <lr_annealing>
        counter: !ref <epoch_counter>

Checkpointer 统一管理 model、scheduler、epoch_counter 三者的状态保存，因此训练可随时中断、随时从最近 checkpoint 恢复；训练结束时也可加载验证损失最小的 checkpoint 做测试。epoch_counter 由 speechbrain/utils/epoch_loop.py 中的 EpochCounter 提供。

默认模型 CustomModel：一个简易 LSTM 语言模型

模板自带 custom_model.py 中的 CustomModel，是一个结构清晰的最小可运行 LSTM LM：

• Embedding：sb.nnet.embedding.Embedding(num_embeddings=output_dim, embedding_dim=embedding_dim)，把 token 索引映射到 256 维稠密向量；
• LSTM：单方向、2 层、隐层 512 的 torch.nn.LSTM；
• 输出层：sb.nnet.linear.Linear(rnn_size, output_dim) + sb.nnet.activations.Softmax(apply_log=True)，输出 log-softmax 概率。

forward 中对 2 维输入自动补时间轴（unsqueeze(dim=1)），这为 beam search 解码阶段逐时间步推理预留了便利；return_hidden=True 时还会返回 LSTM 隐状态，供需要状态传递的解码器使用。

替换为 SpeechBrain 内置语言模型

custom_model.py 的文档字符串明确建议：替换 RNNLM.yaml 中的 !new:custom_model.CustomModel 为内置模型，以获取更强的建模能力：

# RNN 版本
model: !new:speechbrain.lobes.models.RNNLM.RNNLM
    embedding_dim: !ref <emb_dim>
    rnn_size: !ref <rnn_size>
    layers: !ref <layers>
    output_neurons: !ref <output_neurons>   # 需与 tokenizer 词表大小一致

# 或 Transformer 版本
model: !new:speechbrain.lobes.models.transformers.TransformerLM.TransformerLM
    ...

两个内置模型分别位于 speechbrain/lobes/models/RNNLM.py（RNNLM 类）与 speechbrain/lobes/models/transformer/TransformerLM.py（TransformerLM 类）。切换到 TransformerLM 时，还需相应调整 output_neurons（词表大小，本模板对应 1000 个 unigram 单元）以及位置编码、层数等 Transformer 专属超参数。

数据文件与语料格式

模板数据目录 templates/speech_recognition/LM/data 下有三个纯文本文件：

• train.txt（1217 行）：训练转录文本；
• valid.txt：验证文本；
• test.txt：测试文本。

每行是一段大写转录文本，例如 KNIT THREE KNIT TWO TOGETHER KNIT FOUR ...。load_dataset("text", ...) 会按行读取，每行作为一个样本；SpeechBrain 的 tokenizer 将每行编码为 token 序列，再按上文所述拼上 BOS/EOS 形成训练样本。

验证与测试

训练结束后，脚本自动用验证损失最小的 checkpoint 在测试集上评估，输出测试 loss。训练日志中 train_logger（speechbrain.utils.train_logger.FileTrainLogger）会把每个 epoch 的训练/验证 loss 同时写到 results/RNNLM/train_log.txt 和标准输出，便于监控 NewBob 调度是否触发了学习率下降。

若要手动验证推理，可参考 SpeechBrain 的 LM 推理与 ASR 集成（RNNLM 打分器）与 speechbrain/decoders/scorer.py，了解训练好的 LM 如何被 ASR 的 beam search 解码器作为外部先验调用。

小结

本文完整梳理了 SpeechBrain LM 训练模板的配置、代码与运行方式：先跑 Tokenizer 训练出 SentencePiece 模型，再用 python train.py RNNLM.yaml 在 HuggingFace datasets 托管的文本语料上训练 LSTM 语言模型，最终产出可被 ASR 解码器复用的 LM checkpoint。模板通过"小型可复现示例 + 内置可替换模型"的组合，让开发者快速掌握 LM 训练的完整工程链路，并将其无缝嵌入 SpeechBrain 的端到端 ASR 流水线。