Add pred, export and val callbacks (#126)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

ultralytics/yolo/utils/callbacks/base.py

@@ -1,3 +1,7 @@
+# Ultralytics YOLO base callbacks
+
+
+# Trainer callbacks ----------------------------------------------------------------------------------------------------
 def on_pretrain_routine_start(trainer):
     pass
 
@@ -34,26 +38,6 @@ def on_train_epoch_end(trainer):
     pass
 
 
-def on_val_start(trainer):
-    pass
-
-
-def on_val_batch_start(trainer):
-    pass
-
-
-def on_val_image_end(trainer):
-    pass
-
-
-def on_val_batch_end(trainer):
-    pass
-
-
-def on_val_end(trainer):
-    pass
-
-
 def on_fit_epoch_end(trainer):
     pass
 
@@ -74,7 +58,51 @@ def teardown(trainer):
     pass
 
 
+# Validator callbacks --------------------------------------------------------------------------------------------------
+def on_val_start(validator):
+    pass
+
+
+def on_val_batch_start(validator):
+    pass
+
+
+def on_val_batch_end(validator):
+    pass
+
+
+def on_val_end(validator):
+    pass
+
+
+# Predictor callbacks --------------------------------------------------------------------------------------------------
+def on_predict_start(predictor):
+    pass
+
+
+def on_predict_batch_start(predictor):
+    pass
+
+
+def on_predict_batch_end(predictor):
+    pass
+
+
+def on_predict_end(predictor):
+    pass
+
+
+# Exporter callbacks ---------------------------------------------------------------------------------------------------
+def on_export_start(exporter):
+    pass
+
+
+def on_export_end(exporter):
+    pass
+
+
 default_callbacks = {
+    # Run in trainer
     'on_pretrain_routine_start': on_pretrain_routine_start,
     'on_pretrain_routine_end': on_pretrain_routine_end,
     'on_train_start': on_train_start,
@@ -84,16 +112,27 @@ default_callbacks = {
     'on_before_zero_grad': on_before_zero_grad,
     'on_train_batch_end': on_train_batch_end,
     'on_train_epoch_end': on_train_epoch_end,
-    'on_val_start': on_val_start,
-    'on_val_batch_start': on_val_batch_start,
-    'on_val_image_end': on_val_image_end,
-    'on_val_batch_end': on_val_batch_end,
-    'on_val_end': on_val_end,
     'on_fit_epoch_end': on_fit_epoch_end,  # fit = train + val
     'on_model_save': on_model_save,
     'on_train_end': on_train_end,
     'on_params_update': on_params_update,
-    'teardown': teardown}
+    'teardown': teardown,
+
+    # Run in validator
+    'on_val_start': on_val_start,
+    'on_val_batch_start': on_val_batch_start,
+    'on_val_batch_end': on_val_batch_end,
+    'on_val_end': on_val_end,
+
+    # Run in predictor
+    'on_predict_start': on_predict_start,
+    'on_predict_batch_start': on_predict_batch_start,
+    'on_predict_batch_end': on_predict_batch_end,
+    'on_predict_end': on_predict_end,
+
+    # Run in exporter
+    'on_export_start': on_export_start,
+    'on_export_end': on_export_end}
 
 
 def add_integration_callbacks(trainer):

ultralytics/yolo/utils/callbacks/clearml.py

@@ -18,7 +18,7 @@ def _log_images(imgs_dict, group="", step=0):
 
 def on_pretrain_routine_start(trainer):
     # TODO: reuse existing task
-    task = Task.init(project_name=trainer.args.project if trainer.args.project != 'runs/train' else 'YOLOv8',
+    task = Task.init(project_name=trainer.args.project or "YOLOv8",
                      task_name=trainer.args.name,
                      tags=['YOLOv8'],
                      output_uri=True,
@@ -32,7 +32,7 @@ def on_train_epoch_end(trainer):
         _log_images({f.stem: str(f) for f in trainer.save_dir.glob('train_batch*.jpg')}, "Mosaic", trainer.epoch)
 
 
-def on_val_end(trainer):
+def on_fit_epoch_end(trainer):
     if trainer.epoch == 0:
         model_info = {
             "Parameters": get_num_params(trainer.model),
@@ -50,5 +50,5 @@ def on_train_end(trainer):
 callbacks = {
     "on_pretrain_routine_start": on_pretrain_routine_start,
     "on_train_epoch_end": on_train_epoch_end,
-    "on_val_end": on_val_end,
+    "on_fit_epoch_end": on_fit_epoch_end,
     "on_train_end": on_train_end} if clearml else {}

ultralytics/yolo/utils/callbacks/tb.py

@@ -17,11 +17,11 @@ def on_batch_end(trainer):
     _log_scalars(trainer.label_loss_items(trainer.tloss, prefix="train"), trainer.epoch + 1)
 
 
-def on_val_end(trainer):
+def on_fit_epoch_end(trainer):
     _log_scalars(trainer.metrics, trainer.epoch + 1)
 
 
 callbacks = {
     "on_pretrain_routine_start": on_pretrain_routine_start,
-    "on_val_end": on_val_end,
+    "on_fit_epoch_end": on_fit_epoch_end,
     "on_batch_end": on_batch_end}

ultralytics/yolo/utils/callbacks/wb.py

@@ -9,12 +9,11 @@ except (ImportError, AssertionError):
 
 
 def on_pretrain_routine_start(trainer):
-    wandb.init(project=trainer.args.project if trainer.args.project != 'runs/train' else 'YOLOv8',
-               name=trainer.args.name,
-               config=dict(trainer.args)) if not wandb.run else wandb.run
+    wandb.init(project=trainer.args.project or "YOLOv8", name=trainer.args.name, config=dict(
+        trainer.args)) if not wandb.run else wandb.run
 
 
-def on_val_end(trainer):
+def on_fit_epoch_end(trainer):
     wandb.run.log(trainer.metrics, step=trainer.epoch + 1)
     if trainer.epoch == 0:
         model_info = {
@@ -42,5 +41,5 @@ def on_train_end(trainer):
 callbacks = {
     "on_pretrain_routine_start": on_pretrain_routine_start,
     "on_train_epoch_end": on_train_epoch_end,
-    "on_val_end": on_val_end,
+    "on_fit_epoch_end": on_fit_epoch_end,
     "on_train_end": on_train_end} if wandb else {}