k2
k2的愿景是能够将有限状态自动机(FSA)和有限状态转换器(FST)算法无缝集成到基于自动微分的机器学习工具包中,如PyTorch和TensorFlow。对于语音识别应用,这将使得插值和组合各种训练目标(如交叉熵、CTC和MMI)变得容易,并能够联合优化包含多个解码过程的语音识别系统,包括格子重打分和置信度估计。我们希望k2还能有许多其他应用。
我们实现的一个关键算法是通用FSA与"密集"FSA(即对应于神经网络输出的符号对数概率)的剪枝组合。这可以用作ASR解码的快速实现,以及CTC和LF-MMI训练。与现有技术相比,这不会在词错误率方面直接带来优势;但关键是在一个更加通用和可扩展的框架中实现这一点,以便进一步发展ASR技术。
实现
关于我们的实现策略,有几个关键点。
大部分代码是用C++和CUDA编写的。我们实现了一个模板类Ragged
,它很像TensorFlow的RaggedTensor
(实际上我们是独立设计的,后来才被告知TensorFlow使用了相同的想法)。尽管在数据结构层面上有密切的相似性,但设计与TensorFlow和PyTorch有很大不同。大多数情况下,我们不使用简单操作的组合,而是依赖于直接在算法的C++实现中定义的C++11 lambda函数。这些lambda函数中的代码直接操作数据指针,如果后端是CUDA,它们可以为张量的每个元素并行运行。(C++和CUDA代码混合在一起,CUDA内核通过模板实例化)。
不详细查看代码就很难充分描述我们如何使用这些Ragged
对象。由于需要避免顺序处理,这些算法看起来与在CPU上编码的方式非常不同。我们使用的编码模式使计算中最昂贵的部分变得"极易并行化";唯一有些复杂的CUDA操作通常是归约类型的操作,如独占前缀和,我们使用NVidia的cub
库来实现。我们的设计并不特别针对NVidia硬件,我们编写的大部分代码都是相当普通的C++;复杂的CUDA编程主要通过cub库完成,我们用自己方便的接口封装了其中的部分。
然后,有限状态自动机对象被实现为一个针对特定数据类型(表示自动机中的弧的结构体)的Ragged张量。
自动微分
如果你查看现有的代码,你不会发现任何关于自动微分的引用。设计与TensorFlow和PyTorch完全不同(这就是为什么我们没有简单地扩展其中一个工具包)。我们不是从底层开始实现自动微分(通过使单个操作可微分),而是从顶层开始实现,这在这种情况下更加高效(并且往往具有更好的舍入特性)。
举个例子:假设我们正在寻找FSA的最佳路径,并且需要导数。我们通过跟踪输出最佳路径中的每个弧对应于哪个输入弧来实现这一点。(对于更复杂的算法,输出中的一个弧可能对应于一系列输入弧的概率之和)。我们可以在Python层面上使其与PyTorch/TensorFlow的自动微分兼容,例如,在PyTorch中定义一个Function类,该类记住弧之间的这种关系,并执行适当的(稀疏)操作来反向传播权重的导数。
代码的当前状态
我们已经使用pybind11将所有C++代码封装到Python中,并完成了与PyTorch的集成。
我们目前正在编写使用k2的语音识别配方,这些配方托管在一个单独的仓库中。请参见https://github.com/k2-fsa/icefall。
初始发布后的计划
我们目前正在尝试使k2准备好投入生产使用(请参见分支v2.0-pre)。
快速开始
想在不安装任何东西的情况下尝试一下吗?我们已经设置了一个Google Colab。 你可以在https://icefall.readthedocs.io/en/latest/recipes/librispeech/conformer_ctc.html找到更多使用k2进行语音识别的Colab笔记本。