From 52f26e82b3adb9ba89d61e7ca938a00faa7994dd Mon Sep 17 00:00:00 2001 From: Maxim Bonnaerens Date: Fri, 22 Sep 2023 17:36:51 +0200 Subject: [PATCH] Do not materialize mlp_head if no field_heads (#2457) In case it is used in a proposal network we avoid initializing the head. Co-authored-by: Matthew Tancik --- nerfstudio/fields/vanilla_nerf_field.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/nerfstudio/fields/vanilla_nerf_field.py b/nerfstudio/fields/vanilla_nerf_field.py index c9d2bbb78..22e11a25b 100644 --- a/nerfstudio/fields/vanilla_nerf_field.py +++ b/nerfstudio/fields/vanilla_nerf_field.py @@ -74,15 +74,15 @@ class NeRFField(Field): skip_connections=skip_connections, out_activation=nn.ReLU(), ) - - self.mlp_head = MLP( - in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(), - num_layers=head_mlp_num_layers, - layer_width=head_mlp_layer_width, - out_activation=nn.ReLU(), - ) - self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) + + if field_heads: + self.mlp_head = MLP( + in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(), + num_layers=head_mlp_num_layers, + layer_width=head_mlp_layer_width, + out_activation=nn.ReLU(), + ) self.field_heads = nn.ModuleList([field_head() for field_head in field_heads] if field_heads else []) # type: ignore for field_head in self.field_heads: field_head.set_in_dim(self.mlp_head.get_out_dim()) # type: ignore