| 64 | |
| 65 | @registry.register('pointer', 'bahdanau') |
| 66 | class BahdanauPointer(torch.nn.Module): |
| 67 | def __init__(self, query_size, key_size, proj_size): |
| 68 | super().__init__() |
| 69 | self.compute_scores = torch.nn.Sequential( |
| 70 | torch.nn.Linear(query_size + key_size, proj_size), |
| 71 | torch.nn.Tanh(), |
| 72 | torch.nn.Linear(proj_size, 1)) |
| 73 | |
| 74 | def forward(self, query: torch.Tensor, keys: torch.Tensor, attn_mask=None): |
| 75 | # query shape: batch x query_size |
| 76 | # keys shape: batch x num keys x key_size |
| 77 | |
| 78 | # query_expanded shape: batch x num keys x query_size |
| 79 | query_expanded = query.unsqueeze(1).expand(-1, keys.shape[1], -1) |
| 80 | |
| 81 | # scores shape: batch x num keys x 1 |
| 82 | attn_logits = self.compute_scores( |
| 83 | # shape: batch x num keys x query_size + key_size |
| 84 | torch.cat((query_expanded, keys), |
| 85 | dim=2)) |
| 86 | # scores shape: batch x num keys |
| 87 | attn_logits = attn_logits.squeeze(2) |
| 88 | maybe_mask(attn_logits, attn_mask) |
| 89 | return attn_logits |
| 90 | |
| 91 | |
| 92 | @registry.register('attention', 'bahdanau') |