c708aecb18
Signed-off-by: Zhang Jian <jianmusings@gmail.com> Signed-off-by: Zhang <jianmusings@gmail.com> Signed-off-by: Zhang <zhang.jian@u.nus.edu> Signed-off-by: Zhang Jian <e0322744@u.nus.edu> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Zhang Jian <e0322744@u.nus.edu> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
30 lines
861 B
Python
30 lines
861 B
Python
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class GaussianFourierProjection(nn.Module):
|
|
"""Shared Gaussian Fourier features with optional trainable frequencies."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
in_features: int,
|
|
embedding_size: int,
|
|
scale: float = 1.0,
|
|
trainable: bool = True,
|
|
) -> None:
|
|
super().__init__()
|
|
self.in_features = int(in_features)
|
|
self.embedding_size = int(embedding_size)
|
|
weight = torch.randn(self.embedding_size, self.in_features) * scale
|
|
self.weight = nn.Parameter(weight, requires_grad=trainable)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = x.reshape(-1, self.in_features)
|
|
x_proj = 2 * math.pi * x @ self.weight.T
|
|
return torch.cat([x_proj.cos(), x_proj.sin()], dim=-1)
|