| | from dataclasses import dataclass, field |
| | from typing import Optional |
| |
|
| | import torch |
| | import torch.nn as nn |
| | from einops import rearrange |
| |
|
| | from ..utils import BaseModule |
| |
|
| |
|
| | class TriplaneUpsampleNetwork(BaseModule): |
| | @dataclass |
| | class Config(BaseModule.Config): |
| | in_channels: int |
| | out_channels: int |
| |
|
| | cfg: Config |
| |
|
| | def configure(self) -> None: |
| | self.upsample = nn.ConvTranspose2d( |
| | self.cfg.in_channels, self.cfg.out_channels, kernel_size=2, stride=2 |
| | ) |
| |
|
| | def forward(self, triplanes: torch.Tensor) -> torch.Tensor: |
| | triplanes_up = rearrange( |
| | self.upsample( |
| | rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3) |
| | ), |
| | "(B Np) Co Hp Wp -> B Np Co Hp Wp", |
| | Np=3, |
| | ) |
| | return triplanes_up |
| |
|
| |
|
| | class NeRFMLP(BaseModule): |
| | @dataclass |
| | class Config(BaseModule.Config): |
| | in_channels: int |
| | n_neurons: int |
| | n_hidden_layers: int |
| | activation: str = "relu" |
| | bias: bool = True |
| | weight_init: Optional[str] = "kaiming_uniform" |
| | bias_init: Optional[str] = None |
| |
|
| | cfg: Config |
| |
|
| | def configure(self) -> None: |
| | layers = [ |
| | self.make_linear( |
| | self.cfg.in_channels, |
| | self.cfg.n_neurons, |
| | bias=self.cfg.bias, |
| | weight_init=self.cfg.weight_init, |
| | bias_init=self.cfg.bias_init, |
| | ), |
| | self.make_activation(self.cfg.activation), |
| | ] |
| | for i in range(self.cfg.n_hidden_layers - 1): |
| | layers += [ |
| | self.make_linear( |
| | self.cfg.n_neurons, |
| | self.cfg.n_neurons, |
| | bias=self.cfg.bias, |
| | weight_init=self.cfg.weight_init, |
| | bias_init=self.cfg.bias_init, |
| | ), |
| | self.make_activation(self.cfg.activation), |
| | ] |
| | layers += [ |
| | self.make_linear( |
| | self.cfg.n_neurons, |
| | 4, |
| | bias=self.cfg.bias, |
| | weight_init=self.cfg.weight_init, |
| | bias_init=self.cfg.bias_init, |
| | ) |
| | ] |
| | self.layers = nn.Sequential(*layers) |
| |
|
| | def make_linear( |
| | self, |
| | dim_in, |
| | dim_out, |
| | bias=True, |
| | weight_init=None, |
| | bias_init=None, |
| | ): |
| | layer = nn.Linear(dim_in, dim_out, bias=bias) |
| |
|
| | if weight_init is None: |
| | pass |
| | elif weight_init == "kaiming_uniform": |
| | torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu") |
| | else: |
| | raise NotImplementedError |
| |
|
| | if bias: |
| | if bias_init is None: |
| | pass |
| | elif bias_init == "zero": |
| | torch.nn.init.zeros_(layer.bias) |
| | else: |
| | raise NotImplementedError |
| |
|
| | return layer |
| |
|
| | def make_activation(self, activation): |
| | if activation == "relu": |
| | return nn.ReLU(inplace=True) |
| | elif activation == "silu": |
| | return nn.SiLU(inplace=True) |
| | else: |
| | raise NotImplementedError |
| |
|
| | def forward(self, x): |
| | inp_shape = x.shape[:-1] |
| | x = x.reshape(-1, x.shape[-1]) |
| |
|
| | features = self.layers(x) |
| | features = features.reshape(*inp_shape, -1) |
| | out = {"density": features[..., 0:1], "features": features[..., 1:4]} |
| |
|
| | return out |
| |
|
