GPU解码提升40倍,英伟达推进边缘设备部署语音识别,代码已开源
机器之心编译
英伟达近日一篇论文为语音识别技术在边缘设备上的部署带来了福音,其新提出的解码器方法即使在边缘嵌入式 GPU 上也能高效高速地执行。而且这种方法不仅适用于低端硬件,而且也能为数据中心带来显著的效率提升,从而能够识别更多并行的在线音频流。该方法的早期版本已开源。
论文:
代码:
这篇论文提出了一种经过优化的加权式有限状态变换器(WFST/ weighted finite-state transducer)解码器,能够使用图像处理单元(GPU)实现对音频数据的在线流处理和离线批处理。这种解码器能高效利用内存、输入/输出带宽,并为最大化并行使用了一种全新的维特比(Viterbi)实现。内存节省让该解码器能比之前处理更大的图,同时还能支持更多数量的连续流。对 lattice 段进行 GPU 预处理能让中间 lattice 结果在流推理期间返回给请求者。
总体而言,相比于单核 CPU 解码,新提出的改进能实现高达 240 倍的提速,并且解码速度也比当前最佳的 GPU 解码器快 40 倍,同时返回的结果表现相当。从大型数据中心服务器到低功耗边缘设备,该架构可在各种层级的硬件上部署生产级模型。
引言
深度学习研究推动了自动语音识别(ASR)技术的大发展,质量的显著提升让这项技术在许多人机交互用例中得到了实际应用,「流 ASR 即服务(streaming ASR as a service)」的需求也随之水涨船高。通常来说,如要满足这一需求,将需要在数据中心配置大量商用服务器。由于很多用例对时间延迟有很严格的要求,因此人们正在大力投入,意图加快数据中心中模型的推理速度;当然也有研究正致力于实现边缘推理,包括在低功耗设备上实现推理。
在这项研究中,研究者提出了一种全新的加权式有限状态变换器(WFST)实现,其可使用 GPU 和英伟达的 CUDA 编程语言为语音识别任务提供高速解码。他们将该解码器设计成可直接替代现有解码器,而无需修改语言或声学模型。其设计目标是尽力实现最大的灵活性,能支持多路同时音频流的在线识别和 lattice 生成。
研究者还严格限制了该解码器的内存使用,从而可确保 GPU 内存能为大型语言模型和共同常驻的声学模型留有足够的空间。最后,从低功耗嵌入式 GPU 到单个服务器中运行的多个数据中心级 GPU,该算法都能有效地运行。
并行维特比解码
并行式 WFST 解码器通常会遵照串行解码器中的典型操作顺序:对于声学模型(AM)后验的每一帧,该解码器可基于帧值处理发射弧(标签非零的弧),再处理任何非发射弧链,最后执行剪枝。新提出的算法利用了两个类型不同的异步 CUDA 流:一个负责执行计算核,另一个负责执行非阻塞的设备到主机(D2H)lattice token 内存副本。使用第二个用于 D2H 副本的流,无需中止计算流程就能在在线编码期间返回中间结果。
研究者消除了很多常见的面向 CPU 的优化和限制,这种做法有时会妨害并行表现。具体来说,在扩展 token 时不测试新 token 是否唯一。将重复的 token 保留下来留待以后清理对正确性而言是足够的:少量额外的工作能减少对同步和原子操作的依赖。
分批和上下文切换
在 GPU 上,解码核的执行速度很快,其性能受限于核启动的延迟。通过调整解码器的结构,解码器能并行处理多路音频流,通过更长时间地运行这些核能够掩盖启动延迟(由于它们工作负载增加)。
为了同时处理多路音频流,研究者引入了两种不同的机制:干道(channel)和小道(lane)。小道大致等同于神经网络中的批大小,代表了正被解码的话语或流的集合。干道则能为有待继续处理(由于缺少音频或已计算后验)的话语维持状态。这种随时准备 GPU 工作的线程化解码器负责将干道(因为它们已准备就绪)多路复用到小道(因为它们可用)上。这种方案能根据模型和代表性数据与 GPU 的搭配而轻松地调整:增加小道的数量直到收益开始下降,并让通道的数量匹配所测得的吞吐量/xRTF。
内存布局
研究者将这种内存中解码 FST 表示成了一组经过压缩的稀疏行(compressed sparse row,CSR)和附带的元数据,从而可通过直接索引来有效地遍历它们。
给定解码 WFST T = (Σ, Ω, Q, E, ...),其中 Σ 和 Ω 分别为输入和输出标签,Q 为一个有限的状态集,E 为一个有限的过渡集(E_E 是发射过渡),预期内存使用量 M_fst 的计算方式为:
在实践中,这通常等同于用于 FST 的 GPU 内存,大约为磁盘上 FST 大小的 1/3.
解码器所使用的 GPU 内存是有限的,可基于配置好的参数通过一个闭式方程进行计算。下面(2)式给出了全状态解码器的内存足迹(单位是字节),其中包括正被解码的话语和有待进一步解码的话语:
这样,人们就可以基于所需的并行流数量或声学模型/语言模型的大小对该解码器进行扩展。
平衡负载
为了实现并行的最大化,我们最好要能生成大量有差不多同等工作量的线程。文中的做法是在处理每个帧批次时,首先执行一次负载平衡的扩展,其中每个传出弧都由它自己的线程处理,从而生成许多候选 token。然后对这个适应性波束进行调整,并将其用于确定该将哪些候选项加回主队列,以便进一步处理。
另一个不规则的地方源自非发射迭代的缓慢收敛,这会导致小迭代的数量不确定(即长尾)。一旦活动的非发射 token 变得足够低,接下来的迭代就会被一个持续工作的 kernel 处理,直到收敛。在那个持续工作的 kernel 中,每个话语仅拥有一个 CUDA 解码器合作线程阵列(Cooperative Thread Array,CTA),这能提升同步和线程内通信的速度。
一直到解码器中的 lattice 处理阶段,解码器的目标都是发现要为当前帧保留搜索空间的那些子集。基于那个子集构建的后续帧以及在该子集中的任何路径都可能出现在最终 lattice 中。在发现阶段,必须创建和考虑比最终保留的 token 更多的 token(通常多一个数量级)。还有一点,这个发现阶段重在要保持轻量,同时延后任何成本高昂的结构化操作。
为了基于这些 token 生成 lattice,研究者将原始 token 转换成了一种结构化的 CSR 表征。然后,这些数据会被移到主机并用于生成话语末端的最终 lattice。再然后,通过对不能代表其 FST 状态的任何 token 进行「软剪枝」以准备下一帧的 token,具体做法是人工地将它们的外弧度(out-arc degree)归零,然后负载平衡算法可以安全地忽略它们:从而避免指数级增长。
实验和结果
所有的实验都使用了一台英伟达 Tesla V100 GPU 执行,波束=15,lattice-波束=8,最大活动数=10000,除非另有说明。
准确度
速度提升
表 2:离线解码速度(xRTF,波束=15)。
超参数
部署
表 4:当波束大小不同时,在英伟达 GPU 套件上测得的端到端实时吞吐量。
机器之心「SOTA模型」:22大领域、127个任务,机器学习 SOTA 研究一网打尽。