From 800be33cf3fec0dc9dffd80fdddad4e459c87771 Mon Sep 17 00:00:00 2001 From: evonneng Date: Mon, 26 Sep 2022 14:17:04 -0700 Subject: [PATCH] added limit to how many models to save (#546) --- nerfactory/configs/base.py | 2 ++ nerfactory/engine/trainer.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 041f9acf9..1ca2ac4ae 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -146,6 +146,8 @@ class TrainerConfig(PrintableConfig): """whether or not to use mixed precision for training""" relative_model_dir: Path = Path("nerfactory_models/") """relative path to save all checkpoints""" + num_ckpt_to_save: int = 1 + """number of latest checkpoints that we want to save. Default 1 saves only the latest model""" # optional parameters if we want to resume training load_dir: Optional[Path] = None """optionally specify a pre-trained model directory to load from""" diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 18adc3b81..149f00966 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -103,6 +103,7 @@ class Trainer: # directory to save checkpoints self.checkpoint_dir = config.get_checkpoint_dir() logging.info("Saving checkpoints to: %s", self.checkpoint_dir) + self.prev_ckpt_paths = [] # set up viewer if enabled viewer_log_path = self.base_dir / config.viewer.relative_log_filename self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path) @@ -291,6 +292,10 @@ class Trainer: }, ckpt_path, ) + self.prev_ckpt_paths.append(ckpt_path) + if len(self.prev_ckpt_paths) > self.config.trainer.num_ckpt_to_save: + self.prev_ckpt_paths[0].unlink(missing_ok=True) + self.prev_ckpt_paths.pop(0) @profiler.time_function def train_iteration(self, step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: