diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 58785a13f..218050a79 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -131,26 +131,26 @@ class TrainerConfig(PrintableConfig): """Configuration for training regimen""" steps_per_save: int = 1000 - """number of steps between saves""" + """Number of steps between saves.""" steps_per_eval_batch: int = 500 - """number of steps between randomly sampled batches of rays""" + """Number of steps between randomly sampled batches of rays.""" steps_per_eval_image: int = 500 - """number of steps between single eval images""" + """Number of steps between single eval images.""" steps_per_eval_all_images: int = 25000 - """number of steps between eval all images""" + """Number of steps between eval all images.""" max_num_iterations: int = 1000000 - """maximum number of iterations to run""" + """Maximum number of iterations to run.""" mixed_precision: bool = False - """whether or not to use mixed precision for training""" + """Whether or not to use mixed precision for training.""" relative_model_dir: Path = Path("nerfstudio_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""" + """Relative path to save all checkpoints.""" + save_only_latest_checkpoint: bool = True + """Whether to only save the latest checkpoint or all checkpoints.""" # optional parameters if we want to resume training load_dir: Optional[Path] = None - """optionally specify a pre-trained model directory to load from""" + """Optionally specify a pre-trained model directory to load from.""" load_step: Optional[int] = None - """optionally specify model step to load from; if none, will find most recent model in load_dir""" + """Optionally specify model step to load from; if none, will find most recent model in load_dir.""" load_config: Optional[Path] = None diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 4a6546de7..78f7ba0a9 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -103,7 +103,6 @@ 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 ret = None @@ -283,26 +282,28 @@ class Trainer: Args: step: number of steps in training for given checkpoint """ + # possibly make the checkpoint directory if not self.checkpoint_dir.exists(): self.checkpoint_dir.mkdir(parents=True, exist_ok=True) + # save the checkpoint ckpt_path = self.checkpoint_dir / f"step-{step:09d}.ckpt" - if hasattr(self.pipeline, "module"): - pipeline = self.pipeline.module.state_dict() # type: ignore - else: - pipeline = self.pipeline.state_dict() - 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) torch.save( { "step": step, - "pipeline": pipeline, + "pipeline": self.pipeline.module.state_dict() # type: ignore + if hasattr(self.pipeline, "module") + else self.pipeline.state_dict(), "optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()}, "scalers": self.grad_scaler.state_dict(), }, ckpt_path, ) - self.prev_ckpt_paths.append(ckpt_path) + # possibly delete old checkpoints + if self.config.trainer.save_only_latest_checkpoint: + # delete everything else in the checkpoint folder + for f in self.checkpoint_dir.glob("*"): + if f != ckpt_path: + f.unlink() @profiler.time_function def train_iteration(self, step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: