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/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}