diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 08913ce53..fb4e28933 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -264,6 +264,10 @@ class Config(PrintableConfig): self.set_experiment_name() return Path(f"outputs/{self.experiment_name}/{self.method_name}/{self.timestamp}") + def get_checkpoint_dir(self) -> Path: + """Retrieve the checkpoint directory""" + return Path(self.get_base_dir() / self.trainer.relative_model_dir) + def print_to_terminal(self) -> None: """Helper to pretty print config to terminal""" console = Console(width=120) diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 7f65642b5..18adc3b81 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -22,7 +22,6 @@ import functools import logging import os import time -from pathlib import Path from typing import Any, Dict, List, Tuple import torch @@ -102,7 +101,7 @@ class Trainer: self.base_dir = config.get_base_dir() # directory to save checkpoints - self.checkpoint_dir = Path(self.base_dir / self.config.trainer.relative_model_dir) + self.checkpoint_dir = config.get_checkpoint_dir() logging.info("Saving checkpoints to: %s", self.checkpoint_dir) # set up viewer if enabled viewer_log_path = self.base_dir / config.viewer.relative_log_filename @@ -182,7 +181,7 @@ 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: - self._save_checkpoint(step) + self.save_checkpoint(step) writer.write_out_storage() @@ -223,7 +222,7 @@ class Trainer: with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _: num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch try: - self.viewer_state.update_scene(step, self.pipeline.model, num_rays_per_batch) + self.viewer_state.update_scene(self, step, self.pipeline.model, num_rays_per_batch) except RuntimeError: time.sleep(0.03) # sleep to allow buffer to reset assert self.viewer_state.vis is not None @@ -270,7 +269,7 @@ class Trainer: logging.info("No checkpoints to load, training from scratch") @check_main_thread - def _save_checkpoint(self, step: int) -> None: + def save_checkpoint(self, step: int) -> None: """Save the model and optimizers Args: diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index f41b4ed8b..62563c924 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -284,7 +284,7 @@ class ViewerState: # K = camera.get_intrinsics_matrix() # set_persp_intrinsics_matrix(self.vis, K.double().numpy()) - def update_scene(self, step: int, graph: Model, num_rays_per_batch: int) -> None: + def update_scene(self, trainer, step: int, graph: Model, num_rays_per_batch: int) -> None: """updates the scene based on the graph weights Args: @@ -298,6 +298,8 @@ class ViewerState: # check if we should interrupt from a button press? camera_path_payload = self.vis["camera_path_payload"].read() if camera_path_payload: + # save a model checkpoint + trainer.save_checkpoint(step) # write to json file camera_path_filename = camera_path_payload["camera_path_filename"] camera_path = camera_path_payload["camera_path"] diff --git a/scripts/run_eval.py b/scripts/run_eval.py index f5df81c84..e25eafc7a 100755 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -102,7 +102,7 @@ def _eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: # load checkpoints from wherever they were saved # TODO: expose the ability to choose an arbitrary checkpoint - config.trainer.load_dir = config.trainer.model_dir + config.trainer.load_dir = config.get_checkpoint_dir() config.pipeline.datamanager.eval_image_indices = None # setup pipeline (which includes the DataManager)