ultralytics 8.0.59 new MLFlow and feature updates (#1720)

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: St. HeMeow <sheng.heyang@gmail.com>
Co-authored-by: Danny Kim <imbird0312@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Torge Kummerow <CySlider@users.noreply.github.com>
Co-authored-by: dankernel <dkdkernel@gmail.com>
Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com>
Co-authored-by: Roshanlal <roshanlaladchitre103@gmail.com>
Co-authored-by: Lorenzo Mammana <lorenzo.mammana@orobix.com>
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>

ultralytics/yolo/utils/__init__.py

@@ -484,7 +484,7 @@ def get_user_config_dir(sub_dir='Ultralytics'):
     return path
 
 
-USER_CONFIG_DIR = os.getenv('YOLO_CONFIG_DIR', get_user_config_dir())  # Ultralytics settings dir
+USER_CONFIG_DIR = Path(os.getenv('YOLO_CONFIG_DIR', get_user_config_dir()))  # Ultralytics settings dir
 
 
 def emojis(string=''):

ultralytics/yolo/utils/benchmarks.py

@@ -48,8 +48,6 @@ def benchmark(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', imgsz=160, hal
     for i, (name, format, suffix, cpu, gpu) in export_formats().iterrows():  # index, (name, format, suffix, CPU, GPU)
         emoji, filename = '❌', None  # export defaults
         try:
-            if model.task == 'classify':
-                assert i != 11, 'paddle cls exports coming soon'
             assert i != 9 or LINUX, 'Edge TPU export only supported on Linux'
             if i == 10:
                 assert MACOS or LINUX, 'TF.js export only supported on macOS and Linux'

ultralytics/yolo/utils/callbacks/base.py

@@ -147,9 +147,10 @@ def add_integration_callbacks(instance):
     from .clearml import callbacks as clearml_callbacks
     from .comet import callbacks as comet_callbacks
     from .hub import callbacks as hub_callbacks
+    from .mlflow import callbacks as mf_callbacks
     from .tensorboard import callbacks as tb_callbacks
 
-    for x in clearml_callbacks, comet_callbacks, hub_callbacks, tb_callbacks:
+    for x in clearml_callbacks, comet_callbacks, hub_callbacks, tb_callbacks, mf_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/clearml.py

@@ -6,9 +6,9 @@ try:
     import clearml
     from clearml import Task
 
-    assert clearml.__version__  # verify package is not directory
+    assert hasattr(clearml, '__version__')  # verify package is not directory
     assert not TESTS_RUNNING  # do not log pytest
-except (ImportError, AssertionError, AttributeError):
+except (ImportError, AssertionError):
     clearml = None
 
 

ultralytics/yolo/utils/callbacks/comet.py

@@ -6,8 +6,8 @@ try:
     import comet_ml
 
     assert not TESTS_RUNNING  # do not log pytest
-    assert comet_ml.__version__  # verify package is not directory
-except (ImportError, AssertionError, AttributeError):
+    assert hasattr(comet_ml, '__version__')  # verify package is not directory
+except (ImportError, AssertionError):
     comet_ml = None
 
 

ultralytics/yolo/utils/callbacks/mlflow.py

@@ -0,0 +1,75 @@
+# Ultralytics YOLO 🚀, GPL-3.0 license
+
+import os
+import re
+from pathlib import Path
+
+from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING, colorstr
+
+try:
+    import mlflow
+
+    assert not TESTS_RUNNING  # do not log pytest
+    assert hasattr(mlflow, '__version__')  # verify package is not directory
+except (ImportError, AssertionError):
+    mlflow = None
+
+
+def on_pretrain_routine_end(trainer):
+    global mlflow, run, run_id, experiment_name
+
+    if os.environ.get('MLFLOW_TRACKING_URI') is None:
+        mlflow = None
+
+    if mlflow:
+        mlflow_location = os.environ['MLFLOW_TRACKING_URI']  # "http://192.168.xxx.xxx:5000"
+        mlflow.set_tracking_uri(mlflow_location)
+
+        experiment_name = trainer.args.project or 'YOLOv8'
+        experiment = mlflow.get_experiment_by_name(experiment_name)
+        if experiment is None:
+            mlflow.create_experiment(experiment_name)
+        mlflow.set_experiment(experiment_name)
+
+        prefix = colorstr('MLFlow: ')
+        try:
+            run, active_run = mlflow, mlflow.start_run() if mlflow else None
+            if active_run is not None:
+                run_id = active_run.info.run_id
+                LOGGER.info(f'{prefix}Using run_id({run_id}) at {mlflow_location}')
+        except Exception as err:
+            LOGGER.error(f'{prefix}Failing init - {repr(err)}')
+            LOGGER.warning(f'{prefix}Continuing without Mlflow')
+            run = None
+
+        run.log_params(vars(trainer.model.args))
+
+
+def on_fit_epoch_end(trainer):
+    if mlflow:
+        metrics_dict = {f"{re.sub('[()]', '', k)}": float(v) for k, v in trainer.metrics.items()}
+        run.log_metrics(metrics=metrics_dict, step=trainer.epoch)
+
+
+def on_model_save(trainer):
+    if mlflow:
+        run.log_artifact(trainer.last)
+
+
+def on_train_end(trainer):
+    if mlflow:
+        root_dir = Path(__file__).resolve().parents[3]
+        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)},
+                             python_model=run.pyfunc.PythonModel())
+
+
+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/torch_utils.py

@@ -16,10 +16,12 @@ import torch
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
+import torchvision
 
 from ultralytics.yolo.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, __version__
 from ultralytics.yolo.utils.checks import check_version
 
+TORCHVISION_0_10 = check_version(torchvision.__version__, '0.10.0')
 TORCH_1_9 = check_version(torch.__version__, '1.9.0')
 TORCH_1_11 = check_version(torch.__version__, '1.11.0')
 TORCH_1_12 = check_version(torch.__version__, '1.12.0')