Add additional hyperopt callbacks (#1388)

ludwig-ai · Oct 15, 2021 · e6abdf9 · e6abdf9
1 parent 8545f51
commit e6abdf9
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git a/ludwig/callbacks.py b/ludwig/callbacks.py
@@ -25,9 +25,15 @@ def on_cmdline(self, cmd, *args):
     def on_hyperopt_init(self, experiment_name):
         pass
 
+    def on_hyperopt_finish(self, experiment_name):
+        pass
+
     def on_hyperopt_trial_start(self, parameters):
         pass
 
+    def on_hyperopt_trial_end(self, parameters):
+        pass
+
     def on_train_init(
         self,
         base_config,

diff --git a/ludwig/hyperopt/execution.py b/ludwig/hyperopt/execution.py
@@ -1288,6 +1288,10 @@ def run_experiment(
         random_seed=random_seed,
         debug=debug,
     )
+
+    for callback in callbacks or []:
+        callback.on_hyperopt_trial_end(parameters)
+
     return train_stats, eval_stats
 
 

diff --git a/ludwig/hyperopt/run.py b/ludwig/hyperopt/run.py
@@ -337,6 +337,10 @@ def hyperopt(
             save_hyperopt_stats(hyperopt_stats, output_directory)
             logger.info('Hyperopt stats saved to: {}'.format(output_directory))
 
+    if callbacks:
+        for callback in callbacks:
+            callback.on_hyperopt_finish(experiment_name)
+
     logger.info('Finished hyperopt')
 
     return hyperopt_results