19 lines
541 B
Python
19 lines
541 B
Python
import torch
|
|
|
|
|
|
class FourierEmbedder():
|
|
def __init__(self, num_freqs=64, temperature=100):
|
|
|
|
self.num_freqs = num_freqs
|
|
self.temperature = temperature
|
|
self.freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
|
|
|
|
@ torch.no_grad()
|
|
def __call__(self, x, cat_dim=-1):
|
|
"x: arbitrary shape of tensor. dim: cat dim"
|
|
out = []
|
|
for freq in self.freq_bands:
|
|
out.append(torch.sin(freq*x))
|
|
out.append(torch.cos(freq*x))
|
|
return torch.cat(out, cat_dim)
|