Relative positional encoding module (new implementation). Details can be found in https://github.com/espnet/espnet/pull/2816. See : Appendix B in https://arxiv.org/abs/1901.02860 Args: d_model (int): Embedding dimension. dropout_rate (float): Dropout rate. max_
| 713 | |
| 714 | |
| 715 | class EspnetRelPositionalEncoding(torch.nn.Module): |
| 716 | """Relative positional encoding module (new implementation). |
| 717 | |
| 718 | Details can be found in https://github.com/espnet/espnet/pull/2816. |
| 719 | |
| 720 | See : Appendix B in https://arxiv.org/abs/1901.02860 |
| 721 | |
| 722 | Args: |
| 723 | d_model (int): Embedding dimension. |
| 724 | dropout_rate (float): Dropout rate. |
| 725 | max_len (int): Maximum input length. |
| 726 | |
| 727 | """ |
| 728 | |
| 729 | def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000): |
| 730 | """Construct an PositionalEncoding object.""" |
| 731 | super(EspnetRelPositionalEncoding, self).__init__() |
| 732 | self.d_model = d_model |
| 733 | self.xscale = math.sqrt(self.d_model) |
| 734 | self.dropout = torch.nn.Dropout(p=dropout_rate) |
| 735 | self.pe = None |
| 736 | self.extend_pe(torch.tensor(0.0).expand(1, max_len)) |
| 737 | |
| 738 | def extend_pe(self, x: torch.Tensor): |
| 739 | """Reset the positional encodings.""" |
| 740 | if self.pe is not None: |
| 741 | # self.pe contains both positive and negative parts |
| 742 | # the length of self.pe is 2 * input_len - 1 |
| 743 | if self.pe.size(1) >= x.size(1) * 2 - 1: |
| 744 | if self.pe.dtype != x.dtype or self.pe.device != x.device: |
| 745 | self.pe = self.pe.to(dtype=x.dtype, device=x.device) |
| 746 | return |
| 747 | # Suppose `i` means to the position of query vecotr and `j` means the |
| 748 | # position of key vector. We use position relative positions when keys |
| 749 | # are to the left (i>j) and negative relative positions otherwise (i<j). |
| 750 | pe_positive = torch.zeros(x.size(1), self.d_model) |
| 751 | pe_negative = torch.zeros(x.size(1), self.d_model) |
| 752 | position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) |
| 753 | div_term = torch.exp( |
| 754 | torch.arange(0, self.d_model, 2, dtype=torch.float32) |
| 755 | * -(math.log(10000.0) / self.d_model) |
| 756 | ) |
| 757 | pe_positive[:, 0::2] = torch.sin(position * div_term) |
| 758 | pe_positive[:, 1::2] = torch.cos(position * div_term) |
| 759 | pe_negative[:, 0::2] = torch.sin(-1 * position * div_term) |
| 760 | pe_negative[:, 1::2] = torch.cos(-1 * position * div_term) |
| 761 | |
| 762 | # Reserve the order of positive indices and concat both positive and |
| 763 | # negative indices. This is used to support the shifting trick |
| 764 | # as in https://arxiv.org/abs/1901.02860 |
| 765 | pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0) |
| 766 | pe_negative = pe_negative[1:].unsqueeze(0) |
| 767 | pe = torch.cat([pe_positive, pe_negative], dim=1) |
| 768 | self.pe = pe.to(device=x.device, dtype=x.dtype) |
| 769 | |
| 770 | def forward( |
| 771 | self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0 |
| 772 | ) -> Tuple[torch.Tensor, torch.Tensor]: |