| 82 | |
| 83 | class MLPNeRFModel(MetaModule, NeRFModel): |
| 84 | def __init__( |
| 85 | self, |
| 86 | # Positional encoding parameters |
| 87 | n_levels: int = 10, |
| 88 | # MLP parameters |
| 89 | d_hidden: int = 256, |
| 90 | n_density_layers: int = 4, |
| 91 | n_channel_layers: int = 1, |
| 92 | n_channels: int = 3, |
| 93 | sh_degree: int = 4, |
| 94 | activation: str = "relu", |
| 95 | density_activation: str = "exp", |
| 96 | init: Optional[str] = None, |
| 97 | init_scale: float = 1.0, |
| 98 | output_activation: str = "sigmoid", |
| 99 | meta_parameters: bool = False, |
| 100 | trainable_meta: bool = False, |
| 101 | zero_out: bool = True, |
| 102 | register_freqs: bool = True, |
| 103 | posenc_version: str = "v1", |
| 104 | device: torch.device = torch.device("cuda"), |
| 105 | ): |
| 106 | super().__init__() |
| 107 | |
| 108 | # Positional encoding |
| 109 | if register_freqs: |
| 110 | # not used anymore |
| 111 | self.register_buffer( |
| 112 | "freqs", |
| 113 | 2.0 ** torch.arange(n_levels, device=device, dtype=torch.float).view(1, n_levels), |
| 114 | ) |
| 115 | |
| 116 | self.posenc_version = posenc_version |
| 117 | dummy = torch.eye(1, 3) |
| 118 | d_input = encode_position(posenc_version, position=dummy).shape[-1] |
| 119 | |
| 120 | self.n_levels = n_levels |
| 121 | |
| 122 | self.sh_degree = sh_degree |
| 123 | d_sh_coeffs = sh_degree**2 |
| 124 | |
| 125 | self.meta_parameters = meta_parameters |
| 126 | |
| 127 | mlp_cls = ( |
| 128 | partial( |
| 129 | MetaMLP, |
| 130 | meta_scale=False, |
| 131 | meta_shift=False, |
| 132 | meta_proj=True, |
| 133 | meta_bias=True, |
| 134 | trainable_meta=trainable_meta, |
| 135 | ) |
| 136 | if meta_parameters |
| 137 | else MLP |
| 138 | ) |
| 139 | |
| 140 | self.density_mlp = mlp_cls( |
| 141 | d_input=d_input, |