Added at_train_end logic to base pipeline (#1932)

* process additional images

* added on_train_end to base pipeline

* added on_train_end to base pipeline

* remove on_train_end from Base Model and Base Datamanager

* Add new callback type OnTrainEndCallback

* return empty list as default

* add on_train_end to trainingcallbacklocation enum

* make step var into union of int and None

* after train end callback

* Delete unnecessary file

* Delete unnecessary file

* Delete unnecessary file

---------

Co-authored-by: Matthew Tancik <tancik@berkeley.edu>

nerfstudio/engine/callbacks.py

@@ -50,6 +50,7 @@ class TrainingCallbackLocation(Enum):
 
     BEFORE_TRAIN_ITERATION = auto()
     AFTER_TRAIN_ITERATION = auto()
+    AFTER_TRAIN = auto()
 
 
 class TrainingCallback:
@@ -97,6 +98,8 @@ class TrainingCallback:
         elif self.iters is not None:
             if step in self.iters:
                 self.func(*self.args, **self.kwargs, step=step)
+        else:
+            self.func(*self.args, **self.kwargs, step=step)
 
     def run_callback_at_location(self, step: int, location: TrainingCallbackLocation) -> None:
         """Runs the callback if it's supposed to be run at the given location.

nerfstudio/engine/trainer.py

@@ -308,6 +308,10 @@ class Trainer:
         table.add_row("Checkpoint Directory", str(self.checkpoint_dir))
         CONSOLE.print(Panel(table, title="[bold][green]:tada: Training Finished :tada:[/bold]", expand=False))
 
+        # after train end callbacks
+        for callback in self.callbacks:
+            callback.run_callback_at_location(step=step, location=TrainingCallbackLocation.AFTER_TRAIN)
+
         if not self.config.viewer.quit_on_train_completion:
             self._train_complete_viewer()