HUB setup (#108)

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/__init__.py

@@ -249,26 +249,6 @@ def threaded(func):
     return wrapper
 
 
-def get_settings(file=USER_CONFIG_DIR / 'settings.yaml'):
-    """
-    Function that loads a global settings YAML, or creates it and populates it with default values if it does not exist.
-
-    If the datasets or weights directories are set to None, the current working directory will be used.
-    The 'sync' setting determines whether analytics will be synced to help with YOLO development.
-    """
-    from ultralytics.yolo.utils.torch_utils import torch_distributed_zero_first
-
-    with torch_distributed_zero_first(RANK):
-        if not file.exists():
-            settings = {
-                'datasets_dir': None,  # default datasets directory. If None, current working directory is used.
-                'weights_dir': None,  # default weights directory. If None, current working directory is used.
-                'sync': True}  # sync analytics to help with YOLO development
-            yaml_save(file, settings)
-
-    return yaml_load(file)
-
-
 def yaml_save(file='data.yaml', data=None):
     """
     Save YAML data to a file.
@@ -305,6 +285,26 @@ def yaml_load(file='data.yaml'):
         return {**yaml.safe_load(f), 'yaml_file': file}
 
 
+def get_settings(file=USER_CONFIG_DIR / 'settings.yaml'):
+    """
+    Function that loads a global settings YAML, or creates it and populates it with default values if it does not exist.
+
+    If the datasets or weights directories are set to None, the current working directory will be used.
+    The 'sync' setting determines whether analytics will be synced to help with YOLO development.
+    """
+    from ultralytics.yolo.utils.torch_utils import torch_distributed_zero_first
+
+    with torch_distributed_zero_first(RANK):
+        if not file.exists():
+            settings = {
+                'datasets_dir': None,  # default datasets directory. If None, current working directory is used.
+                'weights_dir': None,  # default weights directory. If None, current working directory is used.
+                'sync': True}  # sync analytics to help with YOLO development
+            yaml_save(file, settings)
+
+    return yaml_load(file)
+
+
 # Run below code on utils init -----------------------------------------------------------------------------------------
 
 # Set logger

ultralytics/yolo/utils/callbacks/base.py

@@ -135,11 +135,12 @@ default_callbacks = {
     'on_export_end': on_export_end}
 
 
-def add_integration_callbacks(trainer):
+def add_integration_callbacks(instance):
     from .clearml import callbacks as clearml_callbacks
-    from .tb import callbacks as tb_callbacks
+    from .hub import callbacks as hub_callbacks
+    from .tensorboard import callbacks as tb_callbacks
     from .wb import callbacks as wb_callbacks
 
-    for x in clearml_callbacks, tb_callbacks, wb_callbacks:
+    for x in clearml_callbacks, hub_callbacks, tb_callbacks, wb_callbacks:
         for k, v in x.items():
-            trainer.add_callback(k, v)  # add_callback(name, func)
+            instance.callbacks[k].append(v)  # callback[name].append(func)

ultralytics/yolo/utils/callbacks/hub.py

@@ -0,0 +1,80 @@
+import json
+from time import time
+
+import torch
+
+from ultralytics.hub.utils import PREFIX, sync_analytics
+from ultralytics.yolo.utils import LOGGER
+
+
+def on_pretrain_routine_end(trainer):
+    session = getattr(trainer, 'hub_session', None)
+    if session:
+        # Start timer for upload rate limit
+        LOGGER.info(f"{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀")
+        session.t = {'metrics': time(), 'ckpt': time()}  # start timer on self.rate_limit
+
+
+def on_fit_epoch_end(trainer):
+    session = getattr(trainer, 'hub_session', None)
+    if session:
+        # Upload metrics after val end
+        metrics = trainer.metrics
+        for k, v in metrics.items():
+            if isinstance(v, torch.Tensor):
+                metrics[k] = v.item()
+
+        session.metrics_queue[trainer.epoch] = json.dumps(metrics)  # json string
+        if time() - session.t['metrics'] > session.rate_limits['metrics']:
+            session.upload_metrics()
+            session.t['metrics'] = time()  # reset timer
+            session.metrics_queue = {}  # reset queue
+
+
+def on_model_save(trainer):
+    session = getattr(trainer, 'hub_session', None)
+    if session:
+        # Upload checkpoints with rate limiting
+        is_best = trainer.best_fitness == trainer.fitness
+        if time() - session.t['ckpt'] > session.rate_limits['ckpt']:
+            LOGGER.info(f"{PREFIX}Uploading checkpoint {session.model_id}")
+            session.upload_model(trainer.epoch, trainer.last, is_best)
+            session.t['ckpt'] = time()  # reset timer
+
+
+def on_train_end(trainer):
+    session = getattr(trainer, 'hub_session', None)
+    if session:
+        # Upload final model and metrics with exponential standoff
+        LOGGER.info(f"{PREFIX}Training completed successfully ✅\n"
+                    f"{PREFIX}Uploading final {session.model_id}")
+        session.upload_model(trainer.epoch, trainer.best, map=trainer.metrics['metrics/mAP50(B)'], final=True)
+        session.alive = False  # stop heartbeats
+        LOGGER.info(f"{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀")
+
+
+def on_train_start(trainer):
+    sync_analytics(trainer.args)
+
+
+def on_val_start(validator):
+    sync_analytics(validator.args)
+
+
+def on_predict_start(predictor):
+    sync_analytics(predictor.args)
+
+
+def on_export_start(exporter):
+    sync_analytics(exporter.args)
+
+
+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,
+    "on_train_start": on_train_start,
+    "on_val_start": on_val_start,
+    "on_predict_start": on_predict_start,
+    "on_export_start": on_export_start}