Sample proposal networks on a schedule to speed training (#770)

nerfstudio/configs/method_configs.py

@@ -58,7 +58,7 @@ method_configs["nerfacto"] = Config(
             train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,
             camera_optimizer=CameraOptimizerConfig(
-                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3)
+                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
             ),
         ),
         model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),

nerfstudio/fields/density_fields.py

@@ -62,11 +62,9 @@ class HashMLPDensityField(Field):
         features_per_level=2,
     ) -> None:
         super().__init__()
-
         self.aabb = Parameter(aabb, requires_grad=False)
         self.spatial_distortion = spatial_distortion
         self.use_linear = use_linear
-
         growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
         config = {
@@ -105,12 +103,6 @@ class HashMLPDensityField(Field):
         else:
             positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
         positions_flat = positions.view(-1, 3)
-        # assert all positions are in the range [0, 1]
-        # otherwise print min and max values
-        # assert torch.all(positions >= 0.0) and torch.all(
-        #     positions <= 1.0
-        # ), f"positions: {positions.min()} {positions.max()}"
-
         if not self.use_linear:
             density_before_activation = (
                 self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1).to(positions)

nerfstudio/model_components/ray_samplers.py

@@ -505,12 +505,13 @@ class ProposalNetworkSampler(Sampler):
         num_nerf_samples_per_ray: int = 32,
         num_proposal_network_iterations: int = 2,
         single_jitter: bool = False,
+        update_sched: Callable = lambda x: 1,
     ) -> None:
         super().__init__()
         self.num_proposal_samples_per_ray = num_proposal_samples_per_ray
         self.num_nerf_samples_per_ray = num_nerf_samples_per_ray
         self.num_proposal_network_iterations = num_proposal_network_iterations
-
+        self.update_sched = update_sched
         if self.num_proposal_network_iterations < 1:
             raise ValueError("num_proposal_network_iterations must be >= 1")
 
@@ -519,11 +520,18 @@ class ProposalNetworkSampler(Sampler):
         self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter)
 
         self._anneal = 1.0
+        self._steps_since_update = 0
+        self._step = 0
 
     def set_anneal(self, anneal: float) -> None:
         """Set the anneal value for the proposal network."""
         self._anneal = anneal
 
+    def step_cb(self, step):
+        """Callback to register a training step has passed. This is used to keep track of the sampling schedule"""
+        self._step = step
+        self._steps_since_update += 1
+
     def generate_ray_samples(
         self,
         ray_bundle: Optional[RayBundle] = None,
@@ -538,6 +546,7 @@ class ProposalNetworkSampler(Sampler):
         n = self.num_proposal_network_iterations
         weights = None
         ray_samples = None
+        updated = self._steps_since_update > self.update_sched(self._step) or self._step < 10
         for i_level in range(n + 1):
             is_prop = i_level < n
             num_samples = self.num_proposal_samples_per_ray[i_level] if is_prop else self.num_nerf_samples_per_ray
@@ -551,10 +560,17 @@ 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_fns[i_level](ray_samples.frustums.get_positions())
+                if updated:
+                    # always update on the first step or the inf check in grad scaling crashes
+                    density = density_fns[i_level](ray_samples.frustums.get_positions())
+                else:
+                    with torch.no_grad():
+                        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)
+        if updated:
+            self._steps_since_update = 0
 
         assert ray_samples is not None
         return ray_samples, weights_list, ray_samples_list

nerfstudio/models/nerfacto.py

@@ -61,18 +61,22 @@ class NerfactoModelConfig(ModelConfig):
     """How far along the ray to stop sampling."""
     background_color: Literal["background", "last_sample"] = "last_sample"
     """Whether to randomize the background color."""
-    num_proposal_samples_per_ray: Tuple[int] = (200, 64)
+    num_proposal_samples_per_ray: Tuple[int] = (256, 96)
     """Number of samples per ray for the proposal network."""
-    num_nerf_samples_per_ray: int = 64
+    num_nerf_samples_per_ray: int = 48
     """Number of samples per ray for the nerf network."""
+    proposal_update_every: int = 5
+    """Sample every n steps after the warmup"""
+    proposal_warmup: int = 5000
+    """Scales n from 1 to proposal_update_every over this many steps"""
     num_proposal_iterations: int = 2
     """Number of proposal network iterations."""
     use_same_proposal_network: bool = False
     """Use the same proposal network. Otherwise use different ones."""
     proposal_net_args_list: List[Dict] = field(
         default_factory=lambda: [
-            {"hidden_dim": 16, "log2_hashmap_size": 16, "num_levels": 5, "max_res": 64},
-            {"hidden_dim": 16, "log2_hashmap_size": 16, "num_levels": 5, "max_res": 256},
+            {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 64},
+            {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256},
         ]
     )
     """Arguments for the proposal density fields."""
@@ -129,17 +133,25 @@ class NerfactoModel(Model):
             for i in range(num_prop_nets):
                 prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)]
                 network = HashMLPDensityField(
-                    self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args
+                    self.scene_box.aabb,
+                    spatial_distortion=scene_contraction,
+                    **prop_net_args,
                 )
                 self.proposal_networks.append(network)
             self.density_fns.extend([network.density_fn for network in self.proposal_networks])
 
         # Samplers
+        update_schedule = lambda step: np.clip(
+            np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]),
+            1,
+            self.config.proposal_update_every,
+        )
         self.proposal_sampler = ProposalNetworkSampler(
             num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray,
             num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,
             num_proposal_network_iterations=self.config.num_proposal_iterations,
             single_jitter=self.config.use_single_jitter,
+            update_sched=update_schedule,
         )
 
         # Collider
@@ -186,6 +198,13 @@ class NerfactoModel(Model):
                     func=set_anneal,
                 )
             )
+            callbacks.append(
+                TrainingCallback(
+                    where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
+                    update_every_num_iters=1,
+                    func=self.proposal_sampler.step_cb,
+                )
+            )
         return callbacks
 
     def get_outputs(self, ray_bundle: RayBundle):