ultralytics 8.0.83 Neptune AI logging addition (#2130)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Snyk bot <snyk-bot@snyk.io>
Co-authored-by: Toutatis64 <Toutatis64@users.noreply.github.com>
Co-authored-by: M. Tolga Cangöz <46008593+standardAI@users.noreply.github.com>
Co-authored-by: Talia Bender <85292283+taliabender@users.noreply.github.com>
Co-authored-by: Ophélie Le Mentec <17216799+ouphi@users.noreply.github.com>
Co-authored-by: Kadir Şahin <68073829+ssahinnkadir@users.noreply.github.com>
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>

ultralytics/yolo/utils/callbacks/base.py

@@ -200,11 +200,12 @@ def add_integration_callbacks(instance):
     from .comet import callbacks as comet_callbacks
     from .hub import callbacks as hub_callbacks
     from .mlflow import callbacks as mf_callbacks
+    from .neptune import callbacks as neptune_callbacks
     from .raytune import callbacks as tune_callbacks
     from .tensorboard import callbacks as tb_callbacks
     from .wb import callbacks as wb_callbacks
 
-    for x in clearml_callbacks, comet_callbacks, hub_callbacks, tb_callbacks, mf_callbacks, tune_callbacks, wb_callbacks:
+    for x in clearml_callbacks, comet_callbacks, hub_callbacks, tb_callbacks, mf_callbacks, tune_callbacks, wb_callbacks, neptune_callbacks:
         for k, v in x.items():
             if v not in instance.callbacks[k]:  # prevent duplicate callbacks addition
                 instance.callbacks[k].append(v)  # callback[name].append(func)

ultralytics/yolo/utils/callbacks/mlflow.py

@@ -52,19 +52,12 @@ def on_fit_epoch_end(trainer):
         run.log_metrics(metrics=metrics_dict, step=trainer.epoch)
 
 
-def on_model_save(trainer):
-    """Logs model and metrics to mlflow on save."""
-    if mlflow:
-        run.log_artifact(trainer.last)
-
-
 def on_train_end(trainer):
     """Called at end of train loop to log model artifact info."""
     if mlflow:
         root_dir = Path(__file__).resolve().parents[3]
+        run.log_artifact(trainer.last)
         run.log_artifact(trainer.best)
-        model_uri = f'runs:/{run_id}/'
-        run.register_model(model_uri, experiment_name)
         run.pyfunc.log_model(artifact_path=experiment_name,
                              code_path=[str(root_dir)],
                              artifacts={'model_path': str(trainer.save_dir)},
@@ -74,5 +67,4 @@ def on_train_end(trainer):
 callbacks = {
     'on_pretrain_routine_end': on_pretrain_routine_end,
     'on_fit_epoch_end': on_fit_epoch_end,
-    'on_model_save': on_model_save,
     'on_train_end': on_train_end} if mlflow else {}

ultralytics/yolo/utils/callbacks/neptune.py

@@ -0,0 +1,105 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+import matplotlib.image as mpimg
+import matplotlib.pyplot as plt
+
+from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING
+from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
+
+try:
+    import neptune
+    from neptune.types import File
+
+    assert not TESTS_RUNNING  # do not log pytest
+    assert hasattr(neptune, '__version__')
+except (ImportError, AssertionError):
+    neptune = None
+
+run = None  # NeptuneAI experiment logger instance
+
+
+def _log_scalars(scalars, step=0):
+    """Log scalars to the NeptuneAI experiment logger."""
+    if run:
+        for k, v in scalars.items():
+            run[k].append(value=v, step=step)
+
+
+def _log_images(imgs_dict, group=''):
+    """Log scalars to the NeptuneAI experiment logger."""
+    if run:
+        for k, v in imgs_dict.items():
+            run[f'{group}/{k}'].upload(File(v))
+
+
+def _log_plot(title, plot_path):
+    """Log plots to the NeptuneAI experiment logger."""
+    """
+        Log image as plot in the plot section of NeptuneAI
+
+        arguments:
+        title (str) Title of the plot
+        plot_path (PosixPath or str) Path to the saved image file
+        """
+    img = mpimg.imread(plot_path)
+    fig = plt.figure()
+    ax = fig.add_axes([0, 0, 1, 1], frameon=False, aspect='auto', xticks=[], yticks=[])  # no ticks
+    ax.imshow(img)
+    run[f'Plots/{title}'].upload(fig)
+
+
+def on_pretrain_routine_start(trainer):
+    """Callback function called before the training routine starts."""
+    try:
+        global run
+        run = neptune.init_run(project=trainer.args.project or 'YOLOv8', name=trainer.args.name, tags=['YOLOv8'])
+        run['Configuration/Hyperparameters'] = {k: '' if v is None else v for k, v in vars(trainer.args).items()}
+    except Exception as e:
+        LOGGER.warning(f'WARNING ⚠️ NeptuneAI installed but not initialized correctly, not logging this run. {e}')
+
+
+def on_train_epoch_end(trainer):
+    """Callback function called at end of each training epoch."""
+    _log_scalars(trainer.label_loss_items(trainer.tloss, prefix='train'), trainer.epoch + 1)
+    _log_scalars(trainer.lr, trainer.epoch + 1)
+    if trainer.epoch == 1:
+        _log_images({f.stem: str(f) for f in trainer.save_dir.glob('train_batch*.jpg')}, 'Mosaic')
+
+
+def on_fit_epoch_end(trainer):
+    """Callback function called at end of each fit (train+val) epoch."""
+    if run and trainer.epoch == 0:
+        model_info = {
+            'parameters': get_num_params(trainer.model),
+            'GFLOPs': round(get_flops(trainer.model), 3),
+            'speed(ms)': round(trainer.validator.speed['inference'], 3)}
+        run['Configuration/Model'] = model_info
+    _log_scalars(trainer.metrics, trainer.epoch + 1)
+
+
+def on_val_end(validator):
+    """Callback function called at end of each validation."""
+    if run:
+        # Log val_labels and val_pred
+        _log_images({f.stem: str(f) for f in validator.save_dir.glob('val*.jpg')}, 'Validation')
+
+
+def on_train_end(trainer):
+    """Callback function called at end of training."""
+    if run:
+        # Log final results, CM matrix + PR plots
+        files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))]
+        files = [(trainer.save_dir / f) for f in files if (trainer.save_dir / f).exists()]  # filter
+        for f in files:
+            _log_plot(title=f.stem, plot_path=f)
+        # Log the final model
+        run[f'weights/{trainer.args.name or trainer.args.task}/{str(trainer.best.name)}'].upload(File(str(
+            trainer.best)))
+        run.stop()
+
+
+callbacks = {
+    'on_pretrain_routine_start': on_pretrain_routine_start,
+    'on_train_epoch_end': on_train_epoch_end,
+    'on_fit_epoch_end': on_fit_epoch_end,
+    'on_val_end': on_val_end,
+    'on_train_end': on_train_end} if neptune else {}