Wenet 介绍

时间:2024-02-15 21:06:47
参考资料:
感谢出门问问 的奉献!

1.系统设计

为解决落地问题,同时兼顾简单、高效、易于产品化等的准则,WeNet 做了如下图的三层系统设计。

第一层为 PyTorch 及其生态。

TorchScript 用于开发模型,
Torchaudio 用于 on-the-fly 的特征提取,
DistributedDataParallel 用于分布式训练,
Torch JIT (Just In Time) 用于模型导出,
Pytorch Quantization 用于模型量化,
LibTorch 用于产品模型的推理。

第二层分为研发和产品化两部分。

模型研发阶段,WeNet 使用 TorchScript 做模型开发,以保证实验模型能够正确的导出为产品模型。产品落地阶段,用 LibTorch Production 做模型的推理。

第三层是一个典型的研发模型到落地产品模型的工作流。

其中:
  • Data Prepare:
数据准备部分,WeNet 仅需要准备 kaldi 风格的 wav 列表文件和标注文件。
  • Training:
WeNet 支持 on-the-fly 特征提取、CTC/AED 联合训练和分布式训练。
  • Decoding:
支持 python 环境的模型性能评估,方便在正式部署前的模型调试工作。
  • Export:
研发模型直接导出为产品模型,同时支持导出 float 模型和量化 int8 模型。
  • Runtime:
WeNet 中基于 LibTorch 提供了云端 X86 和嵌入式端 Android 的落地方案,并提供相关工具做准确率、实时率 RTF (real time factor), 延时 (latency) 等产品级别指标的基准测试。

2.模型架构

WeNet 中采用的 U2 模型,如图所示,该模型使用 Joint CTC/AED 的结构,训练时使用 CTC 和 Attention Loss 联合优化,并且通过 dynamic chunk 的训练技巧,使 Shared Encoder 能够处理任意大小的 chunk(即任意长度的语音片段)。在识别的时候,先使用 CTC Decoder 产生得分最高的多个候选结果,再使用 Attention Decoder 对候选结果进行重打分 (Rescoring),并选择重打分后得分最高的结果作为最终识别结果。识别时,当设定 chunk 为无限大的时候,模型需要拿到完整的一句话才能开始做解码,该模型适用于非流式场景,可以充分利用上下文信息获得最佳识别效果;当设定 chunk 为有限大小(如每 0.5 秒语音作为 1 个 chunk)时,Shared Encoder 可以进行增量式的前向运算,同时 CTC Decoder 的结果作为中间结果展示,此时模型可以适用于流时场景,而在流式识别结束时,可利用低延时的重打分算法修复结果,进一步提高最终的识别率。可以看到,依靠该结构,我们同时解决了流式问题和统一模型的问题。

网络结构

Wenet 网络结构设计借鉴 Espnet 的 joint loss 框架,这一框架采取 Conformer Encoder + CTC/attention loss, 利用帧级别的 CTC loss 和 label 级别 attention-based auto-regression loss 联合训练整个网络。 (只看到了encoder的 帧级别ctc loss 和 encoder+decoder 的label级别的 ce loss )

流式 / 非流式的统一结构。

近几年端到端的流式模型和非流式模型往往采用不同结构,一般流式方案可选用 CTC,RNN-T,Monotonic Attention 等方案,在非流式场景下,则通常使 full-attention encoder decoder 模型,牺牲流式特性的而获得更好的识别率。但是,在两种场景下使用两个结构不同的模型会增加模型训练和部署维护的成本,因此最近语音识别的一大研究热点是如何对流式 / 非流式的模型进行统一.
Wenet 针对工业届的具体场景需求,兼顾流式响应速度和识别效果,采用了一套两遍解码的方案,用一个统一的模型结构同时有效的支持流式和非流式场景。
  • CTC 进行第一遍流式解码,该结果可作为流式结果实时返回。
  • 多个候选结果再通过 attention-based decoder 做一遍 teacher forcing rescoring,根据得分重新排序。得到更好的识别结果。

对 Conformer 的流式改进

(encoder部分是confomer结构,要想ctc是流式的,需要对confomer进行流式改进)
原始 Conformer 结构中的 self-attention 需要对整个序列进行 attention 运算,不具备流式特点,同时 convolution 层也依赖于右侧固定长度上下文,且依赖的长度随着模型层数增多而增加。因此,都需要进行改进,保证只依赖有限长度的右侧上下文,从而支持流式推断。
对于 self-attention 层,常见流式改进方法有 chunk 法和 windows 法
  • chunk 法:将序列分成多个 chunk,每帧在 chunk 内部做 attention
  • windows 法:每帧和左右固定长度内的帧做 attention,也叫 look ahead 方法 
从左往右:全序列 attention,左侧 attention,左侧 + chunk 内部分 attention
green means there is a dependency, while white means there is no dependency
Full attention: Full self attention is used in standard Transformer encoder layers, every input at time t depends on the whole inputs,
Left attention: The simplest way to stream it is to make the input t only see itself and the input before t, namely left attention, seeing no right context, as shown in Figure 2 (b), but there is very big degradation compared to full context model. 只能看到左侧及当前帧,来计算attention,性能差。
Chunk Attention: we split the input to several chunks by a fixed chunk size C, the dark green is for the current chunk, for each chunk we have inputs [t+1, t+2, ..., t+C], every chunk depends on itself and the all the previous chunks. Then the whole latency of the encoder depends on the chunk size.
Wenet 选择了 chunk 方案,这种方案可以保证对右侧依赖的长度与网络层数无关。在基于 chunk 的方案中,chunk 的大小会影响着流式最大延时和模型识别率性能,大的 chunk 延时大但性能更好,小的 chunk 延时小但性能会变差。Wenet 使用了一种新的动态 chunk 训练算法,可以使得模型的 chunk 大小动态可变,在运行时,可根据当前场景延时要求和识别率需求手动调整 chunk 大小,而无需重新训练模型。而对于 convolution 层,Wenet 使用一个左侧卷积,以支持流式推断。
(对未来信息做mask,来模仿序列的情况。训练的时候是dynamic chunk,任意大小的chunk都见过了。推理的时候,固定chunk,t+1时刻计算的数量是 t时刻+chunk。chunk 就是多帧一起计算,比如chunk=16,在 t时刻+16帧内,做full attention,在t时刻+16帧之后全部mask,不计算了。)

动态chunk(训练阶段):

流式: 1-25
非流式: max utterance length
下一阶段改进的chunk分布:
非流式方式在每个 Epoch 中占据 50% 的比例
剩下 1-25 (40ms-1s)的 chunk 在剩余的 50% 是等比例分布的,用作流式。
 

解码算法

Wenet 在训练时同时使用了帧级别 deocder 的 CTC loss 和 label 级别 decoder 的 loss。因此在推断时可以使用多种不同的解码方式。目前 Wenet 支持四种解码算法
  • CTC greedy beam search, 帧级别输出,解码过程不合并前缀,最终 n-best 上进行 ctc 序列处理。
(选每一帧概率最高的输出,然后去重去blank,得到最终的序列输出)
  • CTC prefix beam search , 帧级别的解码,合并相同的 ctc 序列前缀。
  • Attention decoder beam search, 基于 cross-attention 的 label 级别解码。
  • CTC + attention rescoring , 将 CTC decoder 的 n-best 结果,通过 attention decoder 进行重打分
CTC prefix beam search+ attention rescoring 方案可以利用 CTC decoder 的输出作为流式结果,在损失很小准确率的情况下支持场景的流式需求,并最终利用 rescoring 改进结果,得到一个识别率更高的最终结果。
考虑一个实时字幕上屏的场景,当演讲者说话时,可利用 CTC decoder 的实时输出结果将字幕实时上屏,让听众立刻看到当前演讲内容的文本,在一句话结束时,会对文本内容进行重打分纠正,最终显示和进行归档的文本是准确率更高的结果。手机 / 车载等语音助理也同样适用该模式。

CTC prefix beam search

每次更新的时候,输出的prefix序列已经是合并过blank和 带blank的重复序列 。每个prefix包含两个得分pb和pnb,该prefix最终的得分为pb+pnb
t 时刻 遍历 beam*beam 个可能,分三种情况,blank、和之前重复、其他情况,分别更新 pb 和 pnb。最后根据这两个概率和选前 beam 个。

Transformer decoder简图:

PS:省略了一些Transformer的基本操作

一个流程示例:

训练: Forward, encoder 选择full attention / dynamic chunk 跑完,根据encoder out 和 label text,计算出帧级别的ctc loss,encoder out 和 label text(shifted right)做decoder,得到的predict text 。和label text,计算label 级的ce。最后loss 按权重相加。
推理:CTC + attention rescoring 的情况:
Encoder 使用指定的chunk size 解码,得到流式的encoder out,使用CTC prefix beam search进行解码,得到beam 个序列hyps及其得分,encoder out 和hyps(shifted right) 去过decoder ,得到 beam个序列对应对得分(b,l,class ,最后一维做softmax,l个概率相加,得到该序列得分),在和ctc得分按权重相加,再排序。得到最好的输出结果。
数据流维度演示
 

3.模型训练

训练脚本

Wenet 提供了 aishell 和 librispeech 的一键 recipe 示例,在 aishell 上可以达到目前业界最好的效果。

特征提取

Wenet 的数据准备非常简单,只需 wav.scp 和 text 文件,即需要训练的音频文件列表和标注文本即可。提取特征在 dataloader 中利用 torchaudio 工具包完成,无需额外的特征提取步骤。
该方案一方面简化了 pipeline,一方面可以在每个 epoch 动态使用不同的 wav augmentatio 方法.
目前 Wenet 中动态特征提取模式仅使用 fbank 特征。 在语音识别中,对于带调的语言如中文可加入 pitch 特征,不过在实验中发现,随着数据量的增大和使用更好的模型结构,往往仅用 fbank 也可达到很高的性能,
另外,Wenet 也支持使用任意特征的方案,但是这种方案需要额外的特征准备步骤,将特征准备成 kaldi 的 ark 格式。
Wenet 还支持一系列可选的 augmentation 算法,如 speed perturbation,spec aug 等方法。

分布式训练

Wenet 目前可以支持在单机单卡和单机多卡上训练,后续会开放多机多卡的功能。

动态 chunk 训练

Wenet 的提出了一套动态 chunk 的训练方法,该方法在训练时,随机为样本选择不同的 chunk 大小,从而允许一个模型可处理不同 chunk 大小的输入,当场景不需要流式识别是,可以使用整句长度作为 chunk 大小,以获得最佳识别性能。当场景需要流式识别时,可根据场景的特点*选择 chunk 大小。
这种方案,在模型训练收敛速度上并不比使用全序列 attention 慢,甚至还可以加速模型的收敛。

4.模型部署

模型导出

Wenet 模型全部由 Torchscript 编写,可以直接导出为支持 libtorch 运行时推断的模型,在运行时,利用 libtorch 导出的接口完成网络前向计算,而特征提取和解码算法则通过 C++ 代码实现.

模型裁剪

目前 Wenet 使用 dynamic 量化来对模型进行压缩,该量化方法对代码结构破坏最小,而且也最为便捷。量化后的模型性能几乎没有损失。

实时率

实时率即处理时间 / 语音时长。更低的实时率意味着可以用更少的计算资源处理更多的语音,即能够为企业带来更低的服务器运营成本。通过模型量化,Wenet 的实时率在服务器端和手机端均可以达到 0.1 以下。即处理 10 秒音频需要的时间不超过 1 秒。
Wenet 的实时率可达 0.1 以内

延时性能

语音识别中存在多种延时需要考虑,如流式识别时用户感知到的实时回显的延时,如一句话说完后完成整句识别的延时。我们考虑三种延时:
L1: 模型结构带来的流式延时
L2: 重打分计算部分带来的延时
L3: 用户说话,到最终重打分后的结果之间的时间间隔。
Wenet 的演示非常低,其不到 0.4s 的流式延时和 0.14s 的整句延时,使得用户在使用时几乎感受不到延时.