Compute 'Scaled Dot Product Attention
(query, key, value, mask=None, dropout=None)
| 95 | |
| 96 | |
| 97 | def attention(query, key, value, mask=None, dropout=None): |
| 98 | "Compute 'Scaled Dot Product Attention'" |
| 99 | d_k = query.size(-1) |
| 100 | scores = torch.matmul(query, key.transpose(-2, -1)) \ |
| 101 | / math.sqrt(d_k) |
| 102 | if mask is not None: |
| 103 | scores = scores.masked_fill(mask == 0, -1e9) |
| 104 | p_attn = F.softmax(scores, dim = -1) |
| 105 | if dropout is not None: |
| 106 | p_attn = dropout(p_attn) |
| 107 | # return torch.matmul(p_attn, value), scores.squeeze(1).squeeze(1) |
| 108 | return torch.matmul(p_attn, value), p_attn |
| 109 | |
| 110 | def sparse_attention(query, key, value, alpha, mask=None, dropout=None): |
| 111 | "Compute 'Scaled Dot Product Attention'" |