|
import math |
|
|
|
import torch |
|
import torch.nn as nn |
|
|
|
|
|
class PhaseModulatedFourierEmbedder(torch.nn.Module): |
|
def __init__( |
|
self, |
|
num_freqs: int, |
|
input_dim: int = 3, |
|
): |
|
""" |
|
Initializes the PhaseModulatedFourierEmbedder class. |
|
Args: |
|
num_freqs (int): The number of frequencies to be used. |
|
input_dim (int, optional): The dimension of the input. Defaults to 3. |
|
Attributes: |
|
weight (torch.nn.Parameter): The weight parameter initialized with random values. |
|
carrier (torch.Tensor): The carrier frequencies calculated based on the Nyquist-Shannon sampling theorem. |
|
out_dim (int): The output dimension calculated based on the input dimension and number of frequencies. |
|
""" |
|
|
|
super().__init__() |
|
|
|
self.weight = nn.Parameter( |
|
torch.randn(input_dim, num_freqs) * math.sqrt(0.5 * num_freqs) |
|
) |
|
|
|
|
|
carrier = (num_freqs / 8) ** torch.linspace(1, 0, num_freqs) |
|
carrier = (carrier + torch.linspace(0, 1, num_freqs)) * 2 * torch.pi |
|
self.register_buffer("carrier", carrier, persistent=False) |
|
|
|
self.out_dim = input_dim * (num_freqs * 2 + 1) |
|
|
|
def forward(self, x): |
|
""" |
|
Perform the forward pass of the embedder model. |
|
Args: |
|
x (torch.Tensor): Input tensor of shape (batch_size, ..., input_dim). |
|
Returns: |
|
torch.Tensor: Output tensor of shape (batch_size, ..., output_dim) where |
|
output_dim = input_dim + 2 * input_dim. |
|
""" |
|
|
|
m = x.float().unsqueeze(-1) |
|
fm = (m * self.weight).view(*x.shape[:-1], -1) |
|
pm = (m * 0.5 * torch.pi + self.carrier).view(*x.shape[:-1], -1) |
|
embedding = torch.cat([x, fm.cos() + pm.cos(), fm.sin() + pm.sin()], dim=-1) |
|
|
|
return embedding |
|
|
