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 another whitening module, move balancer to output.

Browse Source
This commit is contained in:
Daniel Povey 2022-12-07 18:07:56 +08:00
parent 8859177bfa
commit 5f5d02ed0c
1 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1496,7 +1496,12 @@ class NonlinAttentionModule(nn.Module):
min_abs=0.01,
)
self.whiten = Whiten(num_groups=1,
self.whiten1 = Whiten(num_groups=1,
whitening_limit=_whitening_schedule(5.0),
prob=(0.025, 0.25),
grad_scale=0.01)
self.whiten2 = Whiten(num_groups=1,
whitening_limit=_whitening_schedule(5.0),
prob=(0.025, 0.25),
grad_scale=0.01)
@ -1539,10 +1544,11 @@ attn_weights: a Tensor of shape (num_heads, batch_size, seq_len, seq_len)
x = torch.matmul(attn_weights, x)
# now x: (num_heads, batch_size, seq_len, head_dim)
x = x.permute(2, 1, 0, 3).reshape(seq_len, batch_size, -1)
x = self.whiten1(x)
x = self.out_proj(x)
x = self.balancer2(x)
x = self.whiten(x)
x = self.whiten2(x)
return x