| 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 |
|
|
