wenet概述#

wenet處理流程#

wav語音經過一系列前處理之后送入encoder，encoder的輸出會給到ctc decoder和attention decoder。其中ctc decoder是深度優(yōu)先的搜索程序，負責搜索出n段候選預測序列，然后ctc的結構與encoder的結果一起送入attention decoder進行預測序列的打分，選出最好的預測序列，輸出結果

模型結構#

encoder#

encoder是12個conformer模塊堆疊

decoder#

decoder基本結構是transformer，先經過字的自注意力再和mel特征進行交叉注意力，wenet的decoder是雙向decoder，有一個預測序列從左到右的打分，還一個從右到左的打分，每個decoder堆疊3個transformer（取決于訓練時的配置文件，也可以是6個）

模型優(yōu)化#

圖優(yōu)化#

結合netron觀察模型圖結構，考慮以下幾個優(yōu)化方向

刪除冗余算子

算子融合

改變算子執(zhí)行順序

Where(MaskedFill)的優(yōu)化#

這段網絡是在做掩碼操作，即輸入一個掩碼 Tensor對數(shù)據(jù)做Mask，第一個Where把不需要的數(shù)據(jù)設置為-inf。經過Softmax之后這些數(shù)據(jù)已經變成了0，但是后面又增加了一個Where，把相同位置再次設置為了0。

這段結構在網絡中出現(xiàn)了12次，單Where算子耗時30 ms相當于多了一倍的計算時間，可以在編譯階段使用圖優(yōu)化進行消除，減少模型計算量。

耗時從30ms降低到15ms

MatMul的優(yōu)化#

1.這段結構在Decoder中出現(xiàn)6次，且屬于計算集中的attention部分。但是C維度在transpose之后只有8，我們的TPU有64個lane，算力沒有完全利用起來

2.有transpose隔斷了layer group，增加了數(shù)據(jù)搬運

優(yōu)化：

1.可以使用hdim_is_batch優(yōu)化，把attention的head放在h維。為了保證網絡變換前后等效，需要在matmul后面新生成一個transpose。

2.生成新的transpose之后，實現(xiàn)masked fill算子、softmax算子的transpose move down的優(yōu)化pattern，使得tranpose的執(zhí)行順序可以放到該段結構結束處，同時與結束處原本有的tranpose做抵消，達到減少數(shù)據(jù)搬運的目的。

3.由于消除了transpose，使得這段網絡可以做到local layer，同時因為把349放到c維度了，又可以充分利用64個lane的計算資源了

其余的算子經過transpose move down，可以實現(xiàn)transpose的一路下移，在局部網絡中讓C可以保持349，使得64個lane可以獲得更充分的利用

算子層面優(yōu)化#

Where(MaskedFill)的優(yōu)化#

MaskedFill如果全走Global耗時30ms，即便減少一半的算子數(shù)量還是15ms。而Select算子有l(wèi)ocal的實現(xiàn)，同時可以通過參數(shù)配置完成MaskedFill的功能，但不支持廣播。所以在編譯階段加入Tile完成廣播，從而支持Local Layer。

但引入了Tile，Tile操作本身耗時3.8ms，代價可接受，后續(xù)可以進一步優(yōu)化

MaskedFill算子從30ms 減半數(shù)量后到15ms，引入tile之后減少到3.8ms(Tile)+127us(MaskedFill)

后續(xù)考慮使用bdc完成tile操作，完成進一步優(yōu)化

CPU Layer的優(yōu)化#

兩個CPU Gather PT操作占用456ms，可以使用dma的Gather操作在TPU實現(xiàn)算子

Gather PT算子從456ms減少到68us

優(yōu)化結果#