| |
| |
| |
| |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from modules.general.utils import Linear |
|
|
|
|
| class PositionEncoder(nn.Module): |
| r"""Encoder of positional embedding, generates PE and then |
| feed into 2 full-connected layers with ``SiLU``. |
| |
| Args: |
| d_raw_emb: The dimension of raw embedding vectors. |
| d_out: The dimension of output embedding vectors, default to ``d_raw_emb``. |
| d_mlp: The dimension of hidden layer in MLP, default to ``d_raw_emb`` * 4. |
| activation_function: The activation function used in MLP, default to ``SiLU``. |
| n_layer: The number of layers in MLP, default to 2. |
| max_period: controls the minimum frequency of the embeddings. |
| """ |
|
|
| def __init__( |
| self, |
| d_raw_emb: int = 128, |
| d_out: int = None, |
| d_mlp: int = None, |
| activation_function: str = "SiLU", |
| n_layer: int = 2, |
| max_period: int = 10000, |
| ): |
| super().__init__() |
|
|
| self.d_raw_emb = d_raw_emb |
| self.d_out = d_raw_emb if d_out is None else d_out |
| self.d_mlp = d_raw_emb * 4 if d_mlp is None else d_mlp |
| self.n_layer = n_layer |
| self.max_period = max_period |
|
|
| if activation_function.lower() == "silu": |
| self.activation_function = "SiLU" |
| elif activation_function.lower() == "relu": |
| self.activation_function = "ReLU" |
| elif activation_function.lower() == "gelu": |
| self.activation_function = "GELU" |
| else: |
| raise ValueError("activation_function must be one of SiLU, ReLU, GELU") |
| self.activation_function = activation_function |
|
|
| tmp = [Linear(self.d_raw_emb, self.d_mlp), getattr(nn, activation_function)()] |
| for _ in range(self.n_layer - 1): |
| tmp.append(Linear(self.d_mlp, self.d_mlp)) |
| tmp.append(getattr(nn, activation_function)()) |
| tmp.append(Linear(self.d_mlp, self.d_out)) |
|
|
| self.out = nn.Sequential(*tmp) |
|
|
| def forward(self, steps: torch.Tensor) -> torch.Tensor: |
| r"""Create and return sinusoidal timestep embeddings directly. |
| |
| Args: |
| steps: a 1D Tensor of N indices, one per batch element. |
| These may be fractional. |
| |
| Returns: |
| an [N x ``d_out``] Tensor of positional embeddings. |
| """ |
|
|
| half = self.d_raw_emb // 2 |
| freqs = torch.exp( |
| -math.log(self.max_period) |
| / half |
| * torch.arange(half, dtype=torch.float32, device=steps.device) |
| ) |
| args = steps[:, None].float() * freqs[None] |
| embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| if self.d_raw_emb % 2: |
| embedding = torch.cat( |
| [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 |
| ) |
| return self.out(embedding) |
|
|
