语音识别引擎Wenet系列四：进阶话题：Cache「建议收藏」

欢迎大家来到IT世界,在知识的湖畔探索吧!

第一章节可参考

● 第1节: 端到端语音识别基础

CTC目标函数

Attention-based Encoder Decoder

联合建模

神经网络类型

流式语音识别

● 第2节: Wenet中的神经网络设计与实现

Subsampling网络

Encoder Block

模型定义

创建模型

前向计算

其他接口

模型入口 ASRModel

Encoder网络

Attention based Decoder网络

CTC Loss

Attention based Decoder Loss

网络的完整结构

● 第3节: 进阶话题:Mask

Subsampling中的mask

Conformer Block中的Conv的mask

MultiHeadedAttention Module的Mask实现

Chunk-based mask

处理Padding对Loss的影响

处理模型输入Padding

问题1:Batch Padding

问题2: 自回归

问题3: Chunk-Based Model

Encoder中的mask

Decoder中的mask

本文讲解第四章

● 第4节: 进阶话题:Cache

Runtime流式解码

Python流式解码

BaseEncoder.forward_chunk()分析

offset

subsampling内部

subsampling_cache

elayers_output_cache

conformer_cnn_cache

进阶话题:Cache

标准的forward是整个序列进行计算，但是在流式推断时，需要chunk级别的forward，因此需要引入cache的概念，即当前chunk的进行前向计算时，需要拿到上次前向的一些结果作为输入。

什么是cache？

对于流式推断，输入是一个个chunk的到来，对第i个chunk，当计算第k层网络的输出时，由于网络结构存在对左侧上下文的依赖，需要依赖第k-1层网络里在i之前的一些chunks的输出。如果对于当前到来chunk，将其和依赖的chunk序列（比如10层self-attention层，每层依赖左侧4个chunk，则累积起来需要依赖左侧40个chunk）拼起来作为网络输入进行前向，其计算量会比较大。对于那些已经计算过的chunk，可以将那些在计算下一个chunk的输出时需要的中间量保存下来，从而减少重复计算。这种方式就叫cache。

另外，wenet的网络在设计时，对于因果卷积和self-attention的左侧上下文都使用有限长度，因此无论序列多长，每次cache的大小是不变的（不增长）。

仅仅encoder部分涉及chunk计算时的cache。

● 对于CTC decoder，由于是线性层，不需要cache。

● 对于AED decoder，是在计算完整个序列的encoder输出后进行rescoring，不涉及chunk。

Runtime流式解码

asr_model.py中的forward_encoder_chunk()通过jit导出，用于C++ runtime，其内部使用了encoder.py中的forward_chunk()函数。

Python流式解码

如果设置simulate_streaming为True，则会模拟runtime流时解码的过程，将数据分成chunk，依次进行前向计算。该方法的结果，和送入整个序列通过mask进行流式模拟的结果应该是一致的。

forward_chunk_by_chunk()的内部也是使用的forward_chunk()函数。

BaseEncoder.forward_chunk()分析

xs是当前的chunk输入，由于对于单个chunk的前向计算，需要之前的chunk的计算得到的信息，因此这里需要传入相关的三个cache信息。

● subsampling_cache:torch.Tensorsubsampling的输出的cache。即第一个conformer block的输入。

● elayers_output_cache:List[torch.Tensor]第1个到最后1个conformer block的输出的cache。也就是第2个conformer block的输入和CTC层的输入。

● conformer_cnn_cache:List[torch.Tensor]conformer block里的conv层的左侧依赖的输入cache。

cache的大小

● subsampling_cache和elayers_output_cache的大小 由self-attention是对左侧的依赖长度required_cache_size决定。decoding_chunk_size是解码帧级别的chunk大小, num_decoding_left_chunks是self-attention依赖的左侧chunk数。

● conformer_cnn_cache的大小和required_cache_size无关，由casual网络的左侧上下文lorder决定。

函数返回了四个值，包括当前chunk输入对应的输出，更新后的三个cache。

该函数的整个计算过程请参考下图

offset

当按chunk进行输入时，不能直接得到chunk在序列中的位置，需要传入offset给出该chunk在整个序列里的偏移，用于计算positional encoding。

subsampling内部

subsampling内部的计算虽然存在冗余，但是不进行cache。一个是其实现比较复杂，另一个原因是subsampling的计算量占比不大。

subsampling_cache

subsampling的输出的cache。即第一个conformer block的输入。

elayers_output_cache

第1个到最后1个conformer block的输出的cache。也就是第2个conformer block的输入和CTC层的输入。

注意，此处的xs不是当前的chunk，而是当前chunk+cache输入，所以其长度不是chunk_size, 而是chunk_size + required_cache_size。

layer()对应着wenet/transformer/encoder_layer.py中的ConformerEncoderLayer.forward()。下面是其具体过程。

注意，self-attention之前的一些前向计算其实仍然存在冗余，如果对attention层的输入进行cache，而不是对conformer block层的输入cache，可以进一步降低计算量。

conformer_cnn_cache

conformer block里的conv层的左侧依赖的输入cache。

conformer_cnn_cache大小为lorder，即因果卷积左侧依赖。