Reduce memory usage of nerfacto model (#766)

* Reduce memory usage of nerfacto model

* Tweak config

nerfstudio/configs/method_configs.py

@@ -61,7 +61,7 @@ method_configs["nerfacto"] = Config(
                 mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3)
             ),
         ),
-        model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 14),
+        model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
     ),
     optimizers={
         "proposal_networks": {
@@ -73,7 +73,7 @@ method_configs["nerfacto"] = Config(
             "scheduler": None,
         },
     },
-    viewer=ViewerConfig(num_rays_per_chunk=1 << 14),
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
     vis="viewer",
 )
 

nerfstudio/models/nerfacto.py

@@ -200,8 +200,11 @@ class NerfactoModel(Model):
         accumulation = self.renderer_accumulation(weights=weights)
 
         outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth}
-        outputs["weights_list"] = weights_list
-        outputs["ray_samples_list"] = ray_samples_list
+
+        # These use a lot of GPU memory, so we avoid storing them for eval.
+        if self.training:
+            outputs["weights_list"] = weights_list
+            outputs["ray_samples_list"] = ray_samples_list
 
         for i in range(self.config.num_proposal_iterations):
             outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
@@ -212,16 +215,18 @@ class NerfactoModel(Model):
         metrics_dict = {}
         image = batch["image"].to(self.device)
         metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
-        metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
+        if self.training:
+            metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
         return metrics_dict
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = {}
         image = batch["image"].to(self.device)
         loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
-        loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
-            outputs["weights_list"], outputs["ray_samples_list"]
-        )
+        if self.training:
+            loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
+                outputs["weights_list"], outputs["ray_samples_list"]
+            )
         assert metrics_dict is not None and "distortion" in metrics_dict
         loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"]
         return loss_dict

nerfstudio/viewer/server/viewer_utils.py

@@ -639,7 +639,6 @@ class ViewerState:
         check_thread = CheckThread(state=self)
         render_thread = RenderThread(state=self, graph=graph, camera_ray_bundle=camera_ray_bundle)
 
-        oom = False
         with TimeWriter(None, None, write=False) as vis_t:
             check_thread.start()
             render_thread.start()
@@ -647,17 +646,16 @@ class ViewerState:
                 render_thread.join()
                 check_thread.join()
             except IOChangeException:
-                pass
+                del camera_ray_bundle
+                torch.cuda.empty_cache()
             except RuntimeError as e:
-                oom = True
                 self.vis["renderingState/log_errors"].write(
                     "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
                 )
                 print(f"Error: {e}")
-        if oom:
-            del camera_ray_bundle
-            torch.cuda.empty_cache()
-            time.sleep(0.5)  # sleep to allow buffer to reset
+                del camera_ray_bundle
+                torch.cuda.empty_cache()
+                time.sleep(0.5)  # sleep to allow buffer to reset
 
         graph.train()
         outputs = render_thread.vis_outputs