|
import math |
|
|
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
|
|
|
|
class SinusoidalPositionalEncoding(nn.Module): |
|
|
|
def __init__(self, d_model, dropout=0.1, max_len=5000): |
|
super(SinusoidalPositionalEncoding, self).__init__() |
|
self.dropout = nn.Dropout(p=dropout) |
|
|
|
pe = torch.zeros(max_len, d_model) |
|
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) |
|
div_term = torch.arange(0, d_model, 2).float() |
|
div_term = div_term * (-np.log(10000.0) / d_model) |
|
div_term = torch.exp(div_term) |
|
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 = x + self.pe[:x.shape[0]] |
|
return self.dropout(x) |
|
|
|
|
|
class LearnedPositionalEncoding(nn.Module): |
|
|
|
def __init__(self, d_model, dropout=0.1, max_len=5000): |
|
super(LearnedPositionalEncoding, self).__init__() |
|
self.dropout = nn.Dropout(p=dropout) |
|
self.pe = nn.Parameter(torch.randn(max_len, 1, d_model)) |
|
|
|
def forward(self, x): |
|
x = x + self.pe[:x.shape[0]] |
|
return self.dropout(x) |
|
|
|
|
|
def timestep_embedding(timesteps, dim, max_period=10000): |
|
""" |
|
Create sinusoidal timestep embeddings. |
|
:param timesteps: a 1-D Tensor of N indices, one per batch element. |
|
These may be fractional. |
|
:param dim: the dimension of the output. |
|
:param max_period: controls the minimum frequency of the embeddings. |
|
:return: an [N x dim] Tensor of positional embeddings. |
|
""" |
|
half = dim // 2 |
|
idx = torch.arange(start=0, end=half, dtype=torch.float32) |
|
freqs = torch.exp(-math.log(max_period) * idx / |
|
half).to(device=timesteps.device) |
|
args = timesteps[:, None].float() * freqs[None] |
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
|
if dim % 2: |
|
embedding = torch.cat( |
|
[embedding, torch.zeros_like(embedding[:, :1])], dim=-1) |
|
return embedding |
|
|
