(query, key, value, mask=None)
| 9 | return nn.ModuleList([copy.deepcopy(module) for _ in range(N)]) |
| 10 | |
| 11 | def attention(query, key, value, mask=None): |
| 12 | d_k = query.size(-1) |
| 13 | scores = torch.matmul(query, key.transpose(-2, -1).contiguous()) / (d_k ** 0.5) |
| 14 | if mask is not None: |
| 15 | scores = scores.masked_fill(mask == 0, -1e9) |
| 16 | p_attn = F.softmax(scores, dim=-1) |
| 17 | return torch.matmul(p_attn, value), p_attn |
| 18 | |
| 19 | class EncoderDecoder(pl.LightningModule): |
| 20 | def __init__(self, encoder, decoder, src_embed, tgt_embed, generator): |