ultralytics 8.0.48 Edge TPU fix and Metrics updates (#1171)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: majid nasiri <majnasai@gmail.com>

ultralytics/yolo/utils/callbacks/base.py

@@ -151,4 +151,5 @@ def add_integration_callbacks(instance):
 
     for x in clearml_callbacks, comet_callbacks, hub_callbacks, tb_callbacks:
         for k, v in x.items():
-            instance.callbacks[k].append(v)  # callback[name].append(func)
+            if v not in instance.callbacks[k]:  # prevent duplicate callbacks addition
+                instance.callbacks[k].append(v)  # callback[name].append(func)

ultralytics/yolo/utils/callbacks/hub.py

@@ -4,24 +4,33 @@ import json
 from time import time
 
 from ultralytics.hub.utils import PREFIX, traces
-from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING
+from ultralytics.yolo.utils import LOGGER
+from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
 
 
 def on_pretrain_routine_end(trainer):
-    session = not TESTS_RUNNING and getattr(trainer, 'hub_session', None)
+    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
+        session.timers = {'metrics': time(), 'ckpt': time()}  # start timer on session.rate_limit
 
 
 def on_fit_epoch_end(trainer):
     session = getattr(trainer, 'hub_session', None)
     if session:
-        session.metrics_queue[trainer.epoch] = json.dumps(trainer.metrics)  # json string
-        if time() - session.t['metrics'] > session.rate_limits['metrics']:
+        # Upload metrics after val end
+        all_plots = {**trainer.label_loss_items(trainer.tloss, prefix='train'), **trainer.metrics}
+        if trainer.epoch == 0:
+            model_info = {
+                'model/parameters': get_num_params(trainer.model),
+                'model/GFLOPs': round(get_flops(trainer.model), 3),
+                'model/speed(ms)': round(trainer.validator.speed['inference'], 3)}
+            all_plots = {**all_plots, **model_info}
+        session.metrics_queue[trainer.epoch] = json.dumps(all_plots)
+        if time() - session.timers['metrics'] > session.rate_limits['metrics']:
             session.upload_metrics()
-            session.t['metrics'] = time()  # reset timer
+            session.timers['metrics'] = time()  # reset timer
             session.metrics_queue = {}  # reset queue
 
 
@@ -30,21 +39,21 @@ def on_model_save(trainer):
     if session:
         # Upload checkpoints with rate limiting
         is_best = trainer.best_fitness == trainer.fitness
-        if time() - session.t['ckpt'] > session.rate_limits['ckpt']:
+        if time() - session.timers['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
+            session.timers['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-95(B)'], final=True)
-        session.shutdown()  # stop heartbeats
-        LOGGER.info(f'{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀')
+        LOGGER.info(f'{PREFIX}Syncing final model...')
+        session.upload_model(trainer.epoch, trainer.best, map=trainer.metrics.get('metrics/mAP50-95(B)', 0), final=True)
+        session.alive = False  # stop heartbeats
+        LOGGER.info(f'{PREFIX}Done ✅\n'
+                    f'{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀')
 
 
 def on_train_start(trainer):

ultralytics/yolo/utils/callbacks/tensorboard.py

@@ -1,8 +1,12 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
+from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING
 
-from torch.utils.tensorboard import SummaryWriter
+try:
+    from torch.utils.tensorboard import SummaryWriter
 
-from ultralytics.yolo.utils import LOGGER
+    assert not TESTS_RUNNING  # do not log pytest
+except (ImportError, AssertionError):
+    SummaryWriter = None
 
 writer = None  # TensorBoard SummaryWriter instance
 
@@ -18,7 +22,6 @@ def on_pretrain_routine_start(trainer):
     try:
         writer = SummaryWriter(str(trainer.save_dir))
     except Exception as e:
-        writer = None  # TensorBoard SummaryWriter instance
         LOGGER.warning(f'WARNING ⚠️ TensorBoard not initialized correctly, not logging this run. {e}')