From 33d95f3e71e71a2ed72abad9827951e8ce374c65 Mon Sep 17 00:00:00 2001 From: Matias Turkulainen <30566358+maturk@users.noreply.github.com> Date: Wed, 24 May 2023 02:03:55 +0300 Subject: [PATCH] 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 --- nerfstudio/engine/callbacks.py | 3 +++ nerfstudio/engine/trainer.py | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/nerfstudio/engine/callbacks.py b/nerfstudio/engine/callbacks.py index 591711ab5..45aa1d24b 100644 --- a/nerfstudio/engine/callbacks.py +++ b/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. diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 737feb779..4a3a8565c 100644 --- a/nerfstudio/engine/trainer.py +++ b/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()