heimoshuiyu/whisper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

saving the qk matrix in the attention module for convenience

Browse Source
This commit is contained in:
Jong Wook Kim
2022-12-29 23:02:52 -07:00
parent 0b5dcfdef7
commit 68e44bd83c
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
whisper/model.py
View File

@@ -62,6 +62,7 @@ class MultiHeadAttention(nn.Module):
self.key = Linear(n_state, n_state, bias=False) self.key = Linear(n_state, n_state, bias=False)
self.value = Linear(n_state, n_state) self.value = Linear(n_state, n_state)
self.out = Linear(n_state, n_state) self.out = Linear(n_state, n_state)
self.last_qk = None
def forward( def forward(
self, self,
@@ -96,6 +97,8 @@ class MultiHeadAttention(nn.Module):
if mask is not None: if mask is not None:
qk = qk + mask[:n_ctx, :n_ctx] qk = qk + mask[:n_ctx, :n_ctx]
self.last_qk = qk.detach()
w = F.softmax(qk.float(), dim=-1).to(q.dtype) w = F.softmax(qk.float(), dim=-1).to(q.dtype)
return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2) return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2)