Implement the PE function.
| 615 | |
| 616 | |
| 617 | class PositionalEncoding(nn.Module): |
| 618 | "Implement the PE function." |
| 619 | |
| 620 | def __init__(self, dim, dropout, max_len=5000, flipped=False): |
| 621 | super(PositionalEncoding, self).__init__() |
| 622 | self.dropout = nn.Dropout(p=dropout) |
| 623 | self.dim = dim |
| 624 | self.flipped = flipped |
| 625 | |
| 626 | # Compute the positional encodings once in log space. |
| 627 | pe = torch.zeros(max_len, dim) |
| 628 | position = torch.arange(0, max_len).unsqueeze(1) |
| 629 | if self.flipped: |
| 630 | position = -position.flip(dims=[0]) |
| 631 | div_term = torch.exp(torch.arange(0, dim, 2) * -(math.log(10000.0) / dim)) |
| 632 | pe[:, 0::2] = torch.sin(position * div_term) |
| 633 | pe[:, 1::2] = torch.cos(position * div_term) |
| 634 | |
| 635 | pe = pe.unsqueeze(0) |
| 636 | self.register_buffer("pe", pe) |
| 637 | |
| 638 | def forward(self, x): |
| 639 | pe_shape = [1] * (x.ndim - 2) + list(x.shape[-2:-1]) + [self.dim] |
| 640 | if self.flipped: |
| 641 | return self.dropout( |
| 642 | Variable(self.pe[:, -x.size(-2) :].view(pe_shape), requires_grad=False) |
| 643 | ) |
| 644 | else: |
| 645 | return self.dropout( |
| 646 | Variable(self.pe[:, : x.size(-2)].view(pe_shape), requires_grad=False) |
| 647 | ) |
| 648 | |
| 649 | |
| 650 | class PositionalEncodingNd(nn.Module): |
no outgoing calls
no test coverage detected