This module produces sinusoidal positional embeddings of any length. Padding symbols are ignored.
| 86 | |
| 87 | |
| 88 | class SinusoidalPositionalEmbedding(nn.Module): |
| 89 | """This module produces sinusoidal positional embeddings of any length. |
| 90 | |
| 91 | Padding symbols are ignored. |
| 92 | """ |
| 93 | |
| 94 | def __init__(self, embedding_dim, padding_idx, init_size=1024): |
| 95 | super().__init__() |
| 96 | self.embedding_dim = embedding_dim |
| 97 | self.padding_idx = padding_idx |
| 98 | self.weights = SinusoidalPositionalEmbedding.get_embedding( |
| 99 | init_size, |
| 100 | embedding_dim, |
| 101 | padding_idx, |
| 102 | ) |
| 103 | self.register_buffer('_float_tensor', torch.FloatTensor(1)) |
| 104 | |
| 105 | @staticmethod |
| 106 | def get_embedding(num_embeddings, embedding_dim, padding_idx=None): |
| 107 | """Build sinusoidal embeddings. |
| 108 | |
| 109 | This matches the implementation in tensor2tensor, but differs slightly |
| 110 | from the description in Section 3.5 of "Attention Is All You Need". |
| 111 | """ |
| 112 | half_dim = embedding_dim // 2 |
| 113 | emb = math.log(10000) / (half_dim - 1) |
| 114 | emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb) |
| 115 | emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0) |
| 116 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1) |
| 117 | if embedding_dim % 2 == 1: |
| 118 | # zero pad |
| 119 | emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1) |
| 120 | if padding_idx is not None: |
| 121 | emb[padding_idx, :] = 0 |
| 122 | return emb |
| 123 | |
| 124 | def forward(self, input, incremental_state=None, timestep=None, positions=None, **kwargs): |
| 125 | """Input is expected to be of size [bsz x seqlen].""" |
| 126 | bsz, seq_len = input.shape[:2] |
| 127 | max_pos = self.padding_idx + 1 + seq_len |
| 128 | if self.weights is None or max_pos > self.weights.size(0): |
| 129 | # recompute/expand embeddings if needed |
| 130 | self.weights = SinusoidalPositionalEmbedding.get_embedding( |
| 131 | max_pos, |
| 132 | self.embedding_dim, |
| 133 | self.padding_idx, |
| 134 | ) |
| 135 | self.weights = self.weights.to(self._float_tensor) |
| 136 | |
| 137 | if incremental_state is not None: |
| 138 | # positions is the same for every token when decoding a single step |
| 139 | pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len |
| 140 | return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1) |
| 141 | |
| 142 | positions = utils.make_positions(input, self.padding_idx) if positions is None else positions |
| 143 | return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach() |
| 144 | |
| 145 | def max_positions(self): |