diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 69fc4a59a..dbeb1fc86 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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", ) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index b368d1cb4..a76b0534a 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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 diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 893845a3e..b246c0612 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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