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

MultiHeadAttention to return qk as well

Browse Source
This commit is contained in:
Jong Wook Kim
2022-12-30 01:53:06 -07:00
parent 9323b2526c
commit 53807677fe
1 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
whisper/model.py
View File

@@ -82,8 +82,8 @@ class MultiHeadAttention(nn.Module):
k = kv_cache[self.key] k = kv_cache[self.key]
v = kv_cache[self.value] v = kv_cache[self.value]
wv = self.qkv_attention(q, k, v, mask) wv, qk = self.qkv_attention(q, k, v, mask)
return self.out(wv) return self.out(wv), qk
def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None): def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):
n_batch, n_ctx, n_state = q.shape n_batch, n_ctx, n_state = q.shape
@@ -95,9 +95,10 @@ class MultiHeadAttention(nn.Module):
qk = q @ k qk = q @ k
if mask is not None: if mask is not None:
qk = qk + mask[:n_ctx, :n_ctx] qk = qk + mask[:n_ctx, :n_ctx]
qk = qk.float()
w = F.softmax(qk.float(), dim=-1).to(q.dtype) w = F.softmax(qk, 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), qk.detach()
class ResidualAttentionBlock(nn.Module): class ResidualAttentionBlock(nn.Module):
@@ -121,9 +122,9 @@ class ResidualAttentionBlock(nn.Module):
mask: Optional[Tensor] = None, mask: Optional[Tensor] = None,
kv_cache: Optional[dict] = None, kv_cache: Optional[dict] = None,
): ):
x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache) x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache)[0]
if self.cross_attn: if self.cross_attn:
x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache) x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
x = x + self.mlp(self.mlp_ln(x)) x = x + self.mlp(self.mlp_ln(x))
return x return x