Spaces:
Build error
Build error
File size: 2,557 Bytes
e487255 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import math
import torch
from torch import nn
# Protocol for positonal encodings.
# __init__(d_model, max_len=..[, more optionals])
# forward(x: (seq_len, bs, d_model)) -> Tensor of shape (*x.shape[:2],d_model) containing pos. embeddings
class NoPositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=None):
super(NoPositionalEncoding, self).__init__()
pass
def forward(self, x):
return x #* math.sqrt(x.shape[-1])
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0).transpose(0, 1)
self.register_buffer('pe', pe)
def forward(self, x):
x = self.pe[:x.size(0), :] + x # * math.sqrt(x.shape[-1])
return x
class LearnedPositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super(LearnedPositionalEncoding, self).__init__()
self.max_seq_len = max_len
#self.positional_embeddings = nn.Embedding(max_len, d_model)
self.positional_embeddings = nn.Parameter(torch.empty(max_len, d_model))
nn.init.normal_(self.positional_embeddings, mean=0, std=d_model ** -0.5)
def forward(self, x):
seq_len, bs, d_model = x.shape
assert seq_len <= len(self.positional_embeddings), 'seq_len can be at most max_len.'
pos_emb = self.positional_embeddings[:seq_len]
return pos_emb.unsqueeze(1).expand(seq_len, bs, d_model) + x #* math.sqrt(x.shape[-1])
class PairedScrambledPositionalEncodings(LearnedPositionalEncoding):
# TODO check whether it is a problem to use the same perm. for full batch
def forward(self, x):
seq_len, bs, d_model = x.shape
assert seq_len <= len(self.positional_embeddings), 'seq_len can be at most max_len.'
assert len(self.positional_embeddings) % 2 == 0, 'Please specify an even max_len.'
paired_embs = self.positional_embeddings.view(len(self.positional_embeddings), -1, 2)
pos_emb = paired_embs[torch.randperm(len(paired_embs))].view(*self.positional_embeddings.shape)[:seq_len]
return pos_emb.unsqueeze(1).expand(seq_len, bs, d_model) + x #* math.sqrt(x.shape[-1])
|