Improving checkpoint logic (#757)

* working on checkpoint logic

* fixed checkpointing logic

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
 
 

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]]: