option to use different prop networks (#573)

* same or multiple prop networks

* update default

nerfactory/configs/base_configs.py

@@ -97,10 +97,14 @@ base_configs["proposal"] = Config(
         model=ProposalModelConfig(eval_num_rays_per_chunk=8192),
     ),
     optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
         "fields": {
             "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
             "scheduler": None,
-        }
+        },
     },
     viewer=ViewerConfig(enable=True, num_rays_per_chunk=2 << 15),
     logging=LoggingConfig(event_writer="none"),

nerfactory/models/modules/ray_sampler.py

@@ -565,8 +565,10 @@ class ProposalNetworkSampler(Sampler):
     def generate_ray_samples(
         self,
         ray_bundle: Optional[RayBundle] = None,
-        density_fn: Optional[Callable] = None,
+        density_fns: Optional[List[Callable]] = None,
     ) -> Tuple[RaySamples, List, List]:
+        assert ray_bundle is not None
+        assert density_fns is not None
 
         weights_list = []
         ray_samples_list = []
@@ -584,7 +586,7 @@ class ProposalNetworkSampler(Sampler):
                 annealed_weights = torch.pow(weights, self._anneal)
                 ray_samples = self.pdf_sampler(ray_bundle, ray_samples, annealed_weights, num_samples=num_samples)
             if is_prop:
-                density = density_fn(ray_samples.frustums.get_positions())
+                density = density_fns[i_level](ray_samples.frustums.get_positions())
                 weights = ray_samples.get_weights(density)
                 weights_list.append(weights)  # (num_rays, num_samples)
                 ray_samples_list.append(ray_samples)

nerfactory/models/proposal.py

@@ -69,6 +69,8 @@ class ProposalModelConfig(ModelConfig):
     """Number of samples per ray for the nerf network."""
     num_proposal_network_iterations: int = 2
     """Number of proposal network iterations."""
+    use_same_proposal_network: bool = False
+    """Use the same proposal network. Otherwise use different ones."""
     interlevel_loss_mult: float = 1.0
     """Proposal loss multiplier."""
     distortion_loss_mult: float = 0.01
@@ -105,7 +107,18 @@ class ProposalModel(Model):
             )
         else:
             self.field = TCNNInstantNGPField(self.scene_bounds.aabb, spatial_distortion=scene_contraction)
-        self.proposal_network = DensityField(self.scene_bounds.aabb, spatial_distortion=scene_contraction)
+
+        # Build the proposal network(s)
+        if self.config.use_same_proposal_network:
+            network = DensityField(self.scene_bounds.aabb, spatial_distortion=scene_contraction)
+            self.proposal_networks = [network]
+            self.density_fns = [network.density_fn for _ in range(self.config.num_proposal_network_iterations)]
+        else:
+            self.proposal_networks = []
+            for _ in range(self.config.num_proposal_network_iterations):
+                network = DensityField(self.scene_bounds.aabb, spatial_distortion=scene_contraction)
+                self.proposal_networks.append(network)
+            self.density_fns = [network.density_fn for network in self.proposal_networks]
 
         # Collider
         self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
@@ -133,7 +146,11 @@ class ProposalModel(Model):
 
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
         param_groups = {}
-        param_groups["fields"] = list(self.proposal_network.parameters()) + list(self.field.parameters())
+        proposal_networks_params = []
+        for network in self.proposal_networks:
+            proposal_networks_params += list(network.parameters())
+        param_groups["proposal_networks"] = proposal_networks_params
+        param_groups["fields"] = list(self.field.parameters())
         return param_groups
 
     def get_training_callbacks(
@@ -161,9 +178,7 @@ class ProposalModel(Model):
         return callbacks
 
     def get_outputs(self, ray_bundle: RayBundle):
-        ray_samples, weights_list, ray_samples_list = self.proposal_sampler(
-            ray_bundle, density_fn=self.proposal_network.density_fn
-        )
+        ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
         field_outputs = self.field(ray_samples)
         weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
         weights_list.append(weights)
@@ -179,6 +194,7 @@ class ProposalModel(Model):
 
         for i in range(self.config.num_proposal_network_iterations):
             outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
+
         return outputs
 
     def get_metrics_dict(self, outputs, batch):