-
Notifications
You must be signed in to change notification settings - Fork 0
/
positional_encodings.py
70 lines (46 loc) · 2.5 KB
/
positional_encodings.py
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
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])