mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-12-11 06:55:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add Whiten module in NonlinAttentionModule

Browse Source
This commit is contained in:
Daniel Povey 2022-11-21 18:19:23 +08:00
parent 211e3af680
commit 0504f705ec
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
View File

@ -1415,6 +1415,11 @@ class NonlinAttentionModule(nn.Module):
self.in_proj = nn.Linear(channels, 2 * channels, bias=True)
self.whiten = Whiten(num_groups=1,
whitening_limit=10.0,
prob=(0.025, 0.25),
grad_scale=0.01)
self.activation = Identity() # for diagnostics.
self.out_proj = ScaledLinear(channels, channels,
bias=True,
@ -1441,6 +1446,7 @@ attn_weights: a Tensor of shape (num_heads, batch_size, seq_len, seq_len)
# GLU mechanism
x = s.sigmoid() * v
x = self.whiten(x)
(seq_len, batch_size, embed_dim) = x.shape
num_heads = attn_weights.shape[0]