97 lines
3.4 KiB
Python
Executable File
97 lines
3.4 KiB
Python
Executable File
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,
|
|
) -> 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.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.include_input:
|
|
return torch.cat((x, embed.sin(), embed.cos()), dim=-1)
|
|
else:
|
|
return torch.cat((embed.sin(), embed.cos()), dim=-1)
|
|
else:
|
|
return x
|