20 lines
554 B
Python
20 lines
554 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)
|
|
|