18 lines
641 B
Python
18 lines
641 B
Python
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
class PositionalEmbedding(torch.nn.Module):
|
||
|
|
def __init__(self, num_channels, max_positions=10000, endpoint=False):
|
||
|
|
super().__init__()
|
||
|
|
self.num_channels = num_channels
|
||
|
|
self.max_positions = max_positions
|
||
|
|
self.endpoint = endpoint
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
freqs = torch.arange(start=0, end=self.num_channels // 2, dtype=torch.float32, device=x.device)
|
||
|
|
freqs = freqs / (self.num_channels // 2 - (1 if self.endpoint else 0))
|
||
|
|
freqs = (1 / self.max_positions) ** freqs
|
||
|
|
x = x.ger(freqs.to(x.dtype))
|
||
|
|
x = torch.cat([x.cos(), x.sin()], dim=1)
|
||
|
|
return x
|