finish implementation for proposal network, test on prolificdreamer

threestudio-project · Jul 9, 2023 · 0efd25e · 0efd25e
1 parent 68dbae4
commit 0efd25e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/configs/debugging/prolificdreamer-prop.yaml → configs/prolificdreamer-propnet.yaml b/configs/debugging/prolificdreamer-prop.yaml → configs/prolificdreamer-propnet.yaml
@@ -134,9 +134,9 @@ system:
       geometry.encoding:
         lr: 0.01
       geometry.density_network:
-        lr: 0.01
+        lr: 0.001
       geometry.feature_network:
-        lr: 0.01
+        lr: 0.001
       background:
         lr: 0.001
       guidance:

diff --git a/threestudio/models/renderers/nerf_volume_renderer.py b/threestudio/models/renderers/nerf_volume_renderer.py
@@ -378,11 +378,14 @@ def occ_eval_fn(x):
                         step=global_step, occ_eval_fn=occ_eval_fn
                     )
         elif self.cfg.estimator == "proposal":
-            requires_grad = self.proposal_requires_grad_fn(global_step)
-            self.vars_in_forward["requires_grad"] = requires_grad
+            if self.training:
+                requires_grad = self.proposal_requires_grad_fn(global_step)
+                self.vars_in_forward["requires_grad"] = requires_grad
+            else:
+                self.vars_in_forward["requires_grad"] = False
 
     def update_step_end(self, epoch: int, global_step: int) -> None:
-        if self.cfg.estimator == "proposal":
+        if self.cfg.estimator == "proposal" and self.training:
             self.estimator.update_every_n_steps(
                 self.vars_in_forward["trans"],
                 self.vars_in_forward["requires_grad"],