| import math |
| import torch |
| import torch.nn as nn |
|
|
| def make_positions(tensor, padding_idx, left_pad): |
| """Replace non-padding symbols with their position numbers. |
| Position numbers begin at padding_idx+1. |
| Padding symbols are ignored, but it is necessary to specify whether padding |
| is added on the left side (left_pad=True) or right side (left_pad=False). |
| """ |
| max_pos = padding_idx + 1 + tensor.size(1) |
| device = tensor.get_device() |
| buf_name = f'range_buf_{device}' |
| if not hasattr(make_positions, buf_name): |
| setattr(make_positions, buf_name, tensor.new()) |
| setattr(make_positions, buf_name, getattr(make_positions, buf_name).type_as(tensor)) |
| if getattr(make_positions, buf_name).numel() < max_pos: |
| torch.arange(padding_idx + 1, max_pos, out=getattr(make_positions, buf_name)) |
| mask = tensor.ne(padding_idx) |
| positions = getattr(make_positions, buf_name)[:tensor.size(1)].expand_as(tensor) |
| if left_pad: |
| positions = positions - mask.size(1) + mask.long().sum(dim=1).unsqueeze(1) |
| new_tensor = tensor.clone() |
| return new_tensor.masked_scatter_(mask, positions[mask]).long() |
|
|
|
|
| class SinusoidalPositionalEmbedding(nn.Module): |
| """This module produces sinusoidal positional embeddings of any length. |
| Padding symbols are ignored, but it is necessary to specify whether padding |
| is added on the left side (left_pad=True) or right side (left_pad=False). |
| """ |
|
|
| def __init__(self, embedding_dim, padding_idx=0, left_pad=0, init_size=128): |
| super().__init__() |
| self.embedding_dim = embedding_dim |
| self.padding_idx = padding_idx |
| self.left_pad = left_pad |
| self.weights = dict() |
| self.register_buffer('_float_tensor', torch.FloatTensor(1)) |
|
|
| @staticmethod |
| def get_embedding(num_embeddings, embedding_dim, padding_idx=None): |
| """Build sinusoidal embeddings. |
| This matches the implementation in tensor2tensor, but differs slightly |
| from the description in Section 3.5 of "Attention Is All You Need". |
| """ |
| half_dim = embedding_dim // 2 |
| emb = math.log(10000) / (half_dim - 1) |
| emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb) |
| emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0) |
| emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1) |
| if embedding_dim % 2 == 1: |
| |
| emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1) |
| if padding_idx is not None: |
| emb[padding_idx, :] = 0 |
| return emb |
|
|
| def forward(self, input): |
| """Input is expected to be of size [bsz x seqlen].""" |
| bsz, seq_len = input.size() |
| max_pos = self.padding_idx + 1 + seq_len |
| device = input.get_device() |
| if device not in self.weights or max_pos > self.weights[device].size(0): |
| |
| self.weights[device] = SinusoidalPositionalEmbedding.get_embedding( |
| max_pos, |
| self.embedding_dim, |
| self.padding_idx, |
| ) |
| self.weights[device] = self.weights[device].type_as(self._float_tensor).to(input.device) |
| positions = make_positions(input, self.padding_idx, self.left_pad) |
| return self.weights[device].index_select(0, positions.contiguous().view(-1)).view(bsz, seq_len, -1).detach() |
|
|
| def max_positions(self): |
| """Maximum number of supported positions.""" |
| return int(1e5) |
