| from typing import Optional
|
|
|
| from torch import Tensor, nn
|
|
|
| from . import vb_layers_initialize as init
|
|
|
|
|
| class Transition(nn.Module):
|
| """Perform a two-layer MLP."""
|
|
|
| def __init__(
|
| self,
|
| dim: int = 128,
|
| hidden: int = 512,
|
| out_dim: Optional[int] = None,
|
| ) -> None:
|
| """Initialize the TransitionUpdate module.
|
|
|
| Parameters
|
| ----------
|
| dim: int
|
| The dimension of the input, default 128
|
| hidden: int
|
| The dimension of the hidden, default 512
|
| out_dim: Optional[int]
|
| The dimension of the output, default None
|
|
|
| """
|
| super().__init__()
|
| if out_dim is None:
|
| out_dim = dim
|
|
|
| self.norm = nn.LayerNorm(dim, eps=1e-5)
|
| self.fc1 = nn.Linear(dim, hidden, bias=False)
|
| self.fc2 = nn.Linear(dim, hidden, bias=False)
|
| self.fc3 = nn.Linear(hidden, out_dim, bias=False)
|
| self.silu = nn.SiLU()
|
| self.hidden = hidden
|
|
|
| init.bias_init_one_(self.norm.weight)
|
| init.bias_init_zero_(self.norm.bias)
|
|
|
| init.lecun_normal_init_(self.fc1.weight)
|
| init.lecun_normal_init_(self.fc2.weight)
|
| init.final_init_(self.fc3.weight)
|
|
|
| def forward(self, x: Tensor, chunk_size: int = None) -> Tensor:
|
| """Perform a forward pass.
|
|
|
| Parameters
|
| ----------
|
| x: torch.Tensor
|
| The input data of shape (..., D)
|
|
|
| Returns
|
| -------
|
| x: torch.Tensor
|
| The output data of shape (..., D)
|
|
|
| """
|
| x = self.norm(x)
|
|
|
| if chunk_size is None or self.training:
|
| x = self.silu(self.fc1(x)) * self.fc2(x)
|
| x = self.fc3(x)
|
| return x
|
| else:
|
|
|
| for i in range(0, self.hidden, chunk_size):
|
| fc1_slice = self.fc1.weight[i : i + chunk_size, :]
|
| fc2_slice = self.fc2.weight[i : i + chunk_size, :]
|
| fc3_slice = self.fc3.weight[:, i : i + chunk_size]
|
| x_chunk = self.silu((x @ fc1_slice.T)) * (x @ fc2_slice.T)
|
| if i == 0:
|
| x_out = x_chunk @ fc3_slice.T
|
| else:
|
| x_out = x_out + x_chunk @ fc3_slice.T
|
| return x_out
|
|
|
