[论文解读] Pkwrap: a PyTorch Package for LF-MMI Training of Acoustic Models

Srikanth Madikeri, Sibo Tong+4 more|arXiv (Cornell University)|Oct 7, 2020

Speech Recognition and Synthesis参考文献 15被引用 5

一句话总结

Pkwrap 是一个 PyTorch 包，通过封装 Kaldi 的 LF-MMI 框架，实现了端到端的 LF-MMI 语言模型训练，将损失函数暴露为 PyTorch 的自动微分函数，并支持基于对齐和无对齐初始化的训练。它在多个数据集上实现了与 Kaldi 相当的 WER，包括在 Switchboard CH 上使用 4-gram LM 时实现了 5.7% 的相对 WER 改进。

ABSTRACT

We present a simple wrapper that is useful to train acoustic models in PyTorch using Kaldi's LF-MMI training framework. The wrapper, called pkwrap (short form of PyTorch kaldi wrapper), enables the user to utilize the flexibility provided by PyTorch in designing model architectures. It exposes the LF-MMI cost function as an autograd function. Other capabilities of Kaldi have also been ported to PyTorch. This includes the parallel training ability when multi-GPU environments are unavailable and decode with graphs created in Kaldi. The package is available on Github at https://github.com/idiap/pkwrap.

研究动机与目标

• 在保留 Kaldi 最先进的 LF-MMI 训练框架的同时，实现在 PyTorch 中灵活的语音模型设计。
• 支持基于对齐和无对齐初始化的 LF-MMI 训练，以匹配 Kaldi 的训练范式。
• 在多 GPU 和非多 GPU 环境中，为 PyTorch 提供 NG-SGD 并行训练支持。
• 通过使用 Kaldi 兼容的数据和解码图，实现无缝的模型加载和解码。
• 简化研究人员和实践者的基于配方的训练和解码工作流。

提出的方法

• 通过 KaldiChainObjfFunction 类，将 Kaldi 的 LF-MMI 目标函数暴露为 PyTorch 的自动微分函数。
• 通过封装 Kaldi 的 NaturalGradientAffineTransform 和 NGState，实现与 PyTorch 兼容的 NG-SGD 优化器，以支持自然梯度更新。
• 使用 pybind11 绑定关键的 Kaldi C++ 函数，实现在 Kaldi 和 PyTorch 张量之间低开销的内存传输。
• 支持 Kaldi 矩阵与 PyTorch 张量之间的双向转换，以实现模型加载和互操作性。
• 提供配方脚本，通过子进程调用 Kaldi 工具，处理数据准备、特征提取、训练样本生成和解码。
• 集成模型合并与指数调度的学习率衰减，与 Kaldi 的默认行为保持一致。

实验结果

研究问题

• RQ1 PyTorch 封装能否在 LF-MMI 语音模型训练中实现与 Kaldi 相当的 WER 性能？
• RQ2 多大程度上可以将 PyTorch 的模型灵活性与 Kaldi 的稳健 LF-MMI 训练流水线相结合？
• RQ3 Kaldi 的 NG-SGD 优化器在 PyTorch 中集成后，对分布式和非分布式训练的效率如何？
• RQ4 该封装能否在最小用户干预下同时支持基于对齐和无对齐初始化的 LF-MMI 训练？
• RQ5 该包与现有 Kaldi 配方在解码和模型评估方面的互操作性如何？

主要发现

• 在 Minilibrispeech 上，Pkwrap 实现了 19.3% 的 WER，优于 Kaldi 的 20.7%，表明其在小规模数据集上具有竞争力。
• 在 Switchboard 上，使用 4-gram LM 时，Pkwrap 在 CH 子集上实现了 5.7% 的相对 WER 改进，在 SWBD 子集上实现了 4.9% 的改进。
• 在多语言 BABEL 数据集上，Pkwrap 在所有四种语言上的 WER 与 Kaldi 相差不到 1%，在 Swahili 上实现了 3.5% 的相对改进（36.4% vs. 37.7%）。
• 在 Librispeech 100h 数据集上，Pkwrap 在 dev-clean 上实现 5.1% 的 WER，在 test-clean 上实现 5.9%，分别与 Kaldi 的 5.5% 和 6.0% 一致，并在 dev-other 上实现了 5.5% 的相对改进。
• 该包成功支持使用 Adam 优化器和指数学习率衰减进行训练，结果与 Kaldi 的 NG-SGD 基线一致。
• 通过 load_kaldi_models 分支，实现了从 Kaldi 到 PyTorch 的模型加载，支持跨框架的模型复用。