From eaa16e535ae47781eeee01eda280e5a07e6b0850 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Thu, 6 Oct 2022 13:07:37 -0700 Subject: [PATCH] max number of iterations logic (#699) * update max iters * deal with max iters --- nerfstudio/configs/method_configs.py | 8 ++++++-- nerfstudio/engine/trainer.py | 8 +++++++- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index c634a7e98..a51706b76 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -52,7 +52,9 @@ descriptions = {k: convert_markup_to_ansi(v) for k, v in descriptions.items()} method_configs["nerfacto"] = Config( method_name="nerfacto", - trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True), + trainer=TrainerConfig( + steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True + ), pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192 @@ -75,7 +77,9 @@ method_configs["nerfacto"] = Config( method_configs["instant-ngp"] = Config( method_name="instant-ngp", - trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True), + trainer=TrainerConfig( + steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True + ), pipeline=DynamicBatchPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=8192), model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192), diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index d196aec10..2753e4e98 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -184,10 +184,16 @@ class Trainer: self.eval_iteration(step) - if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0: + if step_check(step, self.config.trainer.steps_per_save): self.save_checkpoint(step) writer.write_out_storage() + # save checkpoint at the end of training + self.save_checkpoint(step) + if self.config.is_viewer_enabled(): + while True: + self.viewer_state.vis["renderingState/isTraining"].write(False) + self._update_viewer_state(step) def _check_viewer_warnings(self) -> None: """Helper to print out any warnings regarding the way the viewer/loggers are enabled"""