| ''' |
| ----------------------------------------------------------------------------- |
| Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. |
| |
| NVIDIA CORPORATION and its licensors retain all intellectual property |
| and proprietary rights in and to this software, related documentation |
| and any modifications thereto. Any use, reproduction, disclosure or |
| distribution of this software and related documentation without an express |
| license agreement from NVIDIA CORPORATION is strictly prohibited. |
| ----------------------------------------------------------------------------- |
| ''' |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as torch_F |
|
|
|
|
| class MLPforNeuralSDF(torch.nn.Module): |
|
|
| def __init__(self, layer_dims, skip_connection=[], activ=None, use_layernorm=False, use_weightnorm=False, |
| geometric_init=False, out_bias=0., invert=False): |
| """Initialize a multi-layer perceptron with skip connection. |
| Args: |
| layer_dims: A list of integers representing the number of channels in each layer. |
| skip_connection: A list of integers representing the index of layers to add skip connection. |
| """ |
| super().__init__() |
| self.skip_connection = skip_connection |
| self.use_layernorm = use_layernorm |
| self.linears = torch.nn.ModuleList() |
| if use_layernorm: |
| self.layer_norm = torch.nn.ModuleList() |
| |
| layer_dim_pairs = list(zip(layer_dims[:-1], layer_dims[1:])) |
| for li, (k_in, k_out) in enumerate(layer_dim_pairs): |
| if li in self.skip_connection: |
| k_in += layer_dims[0] |
| linear = torch.nn.Linear(k_in, k_out) |
| if geometric_init: |
| self._geometric_init(linear, k_in, k_out, first=(li == 0), |
| skip_dim=(layer_dims[0] if li in self.skip_connection else 0)) |
| if use_weightnorm: |
| linear = torch.nn.utils.weight_norm(linear) |
| self.linears.append(linear) |
| if use_layernorm and li != len(layer_dim_pairs) - 1: |
| self.layer_norm.append(torch.nn.LayerNorm(k_out)) |
| if li == len(layer_dim_pairs) - 1: |
| self.linears[-1].bias.data.fill_(0.0) |
| |
| self.linear_sdf = torch.nn.Linear(k_in, 1) |
| if geometric_init: |
| self._geometric_init_sdf(self.linear_sdf, k_in, out_bias=out_bias, invert=invert) |
| self.activ = activ or torch_F.relu_ |
|
|
| def forward(self, input, with_sdf=True, with_feat=True): |
| feat = input |
| for li, linear in enumerate(self.linears): |
| if li in self.skip_connection: |
| feat = torch.cat([feat, input], dim=-1) |
| if li != len(self.linears) - 1 or with_feat: |
| feat_pre = linear(feat) |
| if self.use_layernorm: |
| feat_pre = self.layer_norm[li](feat_pre) |
| feat_activ = self.activ(feat_pre) |
| if li == len(self.linears) - 1: |
| out = [self.linear_sdf(feat) if with_sdf else None, |
| feat_activ if with_feat else None] |
| feat = feat_activ |
| return out |
|
|
| def _geometric_init(self, linear, k_in, k_out, first=False, skip_dim=0): |
| torch.nn.init.constant_(linear.bias, 0.0) |
| torch.nn.init.normal_(linear.weight, 0.0, np.sqrt(2 / k_out)) |
| if first: |
| torch.nn.init.constant_(linear.weight[:, 3:], 0.0) |
| if skip_dim: |
| torch.nn.init.constant_(linear.weight[:, -skip_dim:], 0.0) |
|
|
| def _geometric_init_sdf(self, linear, k_in, out_bias=0., invert=False): |
| torch.nn.init.normal_(linear.weight, mean=np.sqrt(np.pi / k_in), std=0.0001) |
| torch.nn.init.constant_(linear.bias, -out_bias) |
| if invert: |
| linear.weight.data *= -1 |
| linear.bias.data *= -1 |
|
|
