| import torch |
| import torch.nn as nn |
|
|
|
|
| class FrequencyPositionalEmbedding(nn.Module): |
| """The sin/cosine positional embedding. Given an input tensor `x` of shape [n_batch, ..., c_dim], it converts |
| each feature dimension of `x[..., i]` into: |
| [ |
| sin(x[..., i]), |
| sin(f_1*x[..., i]), |
| sin(f_2*x[..., i]), |
| ... |
| sin(f_N * x[..., i]), |
| cos(x[..., i]), |
| cos(f_1*x[..., i]), |
| cos(f_2*x[..., i]), |
| ... |
| cos(f_N * x[..., i]), |
| x[..., i] # only present if include_input is True. |
| ], here f_i is the frequency. |
| |
| Denote the space is [0 / num_freqs, 1 / num_freqs, 2 / num_freqs, 3 / num_freqs, ..., (num_freqs - 1) / num_freqs]. |
| If logspace is True, then the frequency f_i is [2^(0 / num_freqs), ..., 2^(i / num_freqs), ...]; |
| Otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1)]. |
| |
| Args: |
| num_freqs (int): the number of frequencies, default is 6; |
| logspace (bool): If logspace is True, then the frequency f_i is [..., 2^(i / num_freqs), ...], |
| otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1)]; |
| input_dim (int): the input dimension, default is 3; |
| include_input (bool): include the input tensor or not, default is True. |
| |
| Attributes: |
| frequencies (torch.Tensor): If logspace is True, then the frequency f_i is [..., 2^(i / num_freqs), ...], |
| otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1); |
| |
| out_dim (int): the embedding size, if include_input is True, it is input_dim * (num_freqs * 2 + 1), |
| otherwise, it is input_dim * num_freqs * 2. |
| |
| """ |
|
|
| def __init__( |
| self, |
| num_freqs: int = 6, |
| logspace: bool = True, |
| input_dim: int = 3, |
| include_input: bool = True, |
| include_pi: bool = True, |
| use_pmpe: bool = False, |
| ) -> None: |
| """The initialization""" |
|
|
| super().__init__() |
|
|
| if logspace: |
| frequencies = 2.0 ** torch.arange(num_freqs, dtype=torch.float32) |
| else: |
| frequencies = torch.linspace( |
| 1.0, 2.0 ** (num_freqs - 1), num_freqs, dtype=torch.float32 |
| ) |
|
|
| if include_pi: |
| frequencies *= torch.pi |
|
|
| self.register_buffer("frequencies", frequencies, persistent=False) |
| self.include_input = include_input |
| self.num_freqs = num_freqs |
| self.use_pmpe = use_pmpe |
| if use_pmpe: |
| phase = torch.arange(num_freqs, dtype=torch.float32) |
| for i in range(num_freqs): |
| phase[i] = torch.pow(torch.tensor(num_freqs), 1.0-(i+1)/num_freqs)+(i+1)/num_freqs |
| phase *= torch.pi*2 |
| self.register_buffer("phase", phase, persistent=False) |
|
|
| self.out_dim = self.get_dims(input_dim) |
|
|
| def get_dims(self, input_dim): |
| temp = 1 if self.include_input or self.num_freqs == 0 else 0 |
| out_dim = input_dim * (self.num_freqs * 2 + temp) |
|
|
| return out_dim |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """Forward process. |
| |
| Args: |
| x: tensor of shape [..., dim] |
| |
| Returns: |
| embedding: an embedding of `x` of shape [..., dim * (num_freqs * 2 + temp)] |
| where temp is 1 if include_input is True and 0 otherwise. |
| """ |
|
|
| if self.num_freqs > 0: |
| embed = (x[..., None].contiguous() * self.frequencies).view( |
| *x.shape[:-1], -1 |
| ) |
| if self.use_pmpe: |
| phase = (x[..., None].contiguous()*torch.pi*0.5 + self.phase).view( |
| *x.shape[:-1], -1 |
| ) |
| res = torch.cat((embed.sin()+phase.sin(), embed.cos()+phase.cos()), dim=-1) |
| else: |
| res = torch.cat((embed.sin(), embed.cos()), dim=-1) |
| if self.include_input: |
| return torch.cat((x, res), dim=-1) |
| else: |
| return res |
| else: |
| return x |
|
|
