Pkwrap: a PyTorch Package for LF-MMI Training of Acoustic Models

TL;DR

Pkwrap is a PyTorch package that simplifies LF-MMI training of acoustic models by integrating Kaldi's framework, enabling flexible model design and multi-GPU training.

Contribution

It introduces a PyTorch wrapper for Kaldi's LF-MMI training, combining Kaldi's capabilities with PyTorch's flexibility for acoustic model training.

Findings

Enables LF-MMI training in PyTorch with Kaldi's framework.

Supports multi-GPU training and decoding in PyTorch.

Provides an easy-to-use interface for acoustic model development.

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.