| """Siren MLP https://www.vincentsitzmann.com/siren/""" | |
| from typing import Optional | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| class SineLayer(nn.Module): | |
| """ | |
| Sine layer for the SIREN network. | |
| """ | |
| def __init__( | |
| self, in_features, out_features, bias=True, is_first=False, omega_0=30.0 | |
| ): | |
| super().__init__() | |
| self.omega_0 = omega_0 | |
| self.is_first = is_first | |
| self.in_features = in_features | |
| self.linear = nn.Linear(in_features, out_features, bias=bias) | |
| self.init_weights() | |
| def init_weights(self): | |
| with torch.no_grad(): | |
| if self.is_first: | |
| self.linear.weight.uniform_(-1 / self.in_features, 1 / self.in_features) | |
| else: | |
| self.linear.weight.uniform_( | |
| -np.sqrt(6 / self.in_features) / self.omega_0, | |
| np.sqrt(6 / self.in_features) / self.omega_0, | |
| ) | |
| def forward(self, x): | |
| return torch.sin(self.omega_0 * self.linear(x)) | |
| class Siren(nn.Module): | |
| """Siren network. | |
| Args: | |
| in_dim: Input layer dimension | |
| num_layers: Number of network layers | |
| layer_width: Width of each MLP layer | |
| out_dim: Output layer dimension. Uses layer_width if None. | |
| activation: intermediate layer activation function. | |
| out_activation: output activation function. | |
| """ | |
| def __init__( | |
| self, | |
| in_dim: int, | |
| hidden_layers: int, | |
| hidden_features: int, | |
| out_dim: Optional[int] = None, | |
| outermost_linear: bool = False, | |
| first_omega_0: float = 30, | |
| hidden_omega_0: float = 30, | |
| out_activation: Optional[nn.Module] = None, | |
| ) -> None: | |
| super().__init__() | |
| self.in_dim = in_dim | |
| assert self.in_dim > 0 | |
| self.out_dim = out_dim if out_dim is not None else hidden_features | |
| self.outermost_linear = outermost_linear | |
| self.first_omega_0 = first_omega_0 | |
| self.hidden_omega_0 = hidden_omega_0 | |
| self.hidden_layers = hidden_layers | |
| self.layer_width = hidden_features | |
| self.out_activation = out_activation | |
| self.net = [] | |
| self.net.append( | |
| SineLayer(in_dim, hidden_features, is_first=True, omega_0=first_omega_0) | |
| ) | |
| for _ in range(hidden_layers): | |
| self.net.append( | |
| SineLayer( | |
| hidden_features, | |
| hidden_features, | |
| is_first=False, | |
| omega_0=hidden_omega_0, | |
| ) | |
| ) | |
| if outermost_linear: | |
| final_layer = nn.Linear(hidden_features, self.out_dim) | |
| with torch.no_grad(): | |
| final_layer.weight.uniform_( | |
| -np.sqrt(6 / hidden_features) / hidden_omega_0, | |
| np.sqrt(6 / hidden_features) / hidden_omega_0, | |
| ) | |
| self.net.append(final_layer) | |
| else: | |
| self.net.append( | |
| SineLayer( | |
| hidden_features, | |
| self.out_dim, | |
| is_first=False, | |
| omega_0=hidden_omega_0, | |
| ) | |
| ) | |
| if self.out_activation is not None: | |
| self.net.append(self.out_activation) | |
| self.net = nn.Sequential(*self.net) | |
| def forward(self, model_input): | |
| """Forward pass through the network""" | |
| output = self.net(model_input) | |
| return output | |