Take in model size and number of heads.
(self, h, d_model, dropout=0.1)
| 25 | |
| 26 | class MultiHeadedAttention(torch.nn.Module): |
| 27 | def __init__(self, h, d_model, dropout=0.1): |
| 28 | "Take in model size and number of heads." |
| 29 | super(MultiHeadedAttention, self).__init__() |
| 30 | assert d_model % h == 0 |
| 31 | # We assume d_v always equals d_k |
| 32 | self.d_k = d_model // h |
| 33 | self.h = h |
| 34 | self.linears = clones(torch.nn.Linear(d_model, d_model), 4) |
| 35 | self.attn = None |
| 36 | self.dropout = torch.nn.Dropout(p=dropout) |
| 37 | |
| 38 | def forward(self, query, key, value, mask=None): |
| 39 | "Implements Figure 2" |