Add Comet integration (#153)

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

@@ -137,10 +137,11 @@ default_callbacks = {
 
 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 .tensorboard import callbacks as tb_callbacks
     from .wb import callbacks as wb_callbacks
 
-    for x in clearml_callbacks, hub_callbacks, tb_callbacks, wb_callbacks:
+    for x in clearml_callbacks, comet_callbacks, hub_callbacks, tb_callbacks, wb_callbacks:
         for k, v in x.items():
             instance.callbacks[k].append(v)  # callback[name].append(func)

ultralytics/yolo/utils/callbacks/comet.py

@@ -0,0 +1,43 @@
+from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
+
+try:
+    import comet_ml
+
+except (ModuleNotFoundError, ImportError):
+    comet_ml = None
+
+
+def on_pretrain_routine_start(trainer):
+    experiment = comet_ml.Experiment(project_name=trainer.args.project or "YOLOv8",)
+    experiment.log_parameters(dict(trainer.args))
+
+
+def on_train_epoch_end(trainer):
+    experiment = comet_ml.get_global_experiment()
+    experiment.log_metrics(trainer.label_loss_items(trainer.tloss, prefix="train"), step=trainer.epoch + 1)
+    if trainer.epoch == 1:
+        for f in trainer.save_dir.glob('train_batch*.jpg'):
+            experiment.log_image(f, name=f.stem, step=trainer.epoch + 1)
+
+
+def on_fit_epoch_end(trainer):
+    experiment = comet_ml.get_global_experiment()
+    experiment.log_metrics(trainer.metrics, step=trainer.epoch + 1)
+    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[1], 3)}
+        experiment.log_metrics(model_info, step=trainer.epoch + 1)
+
+
+def on_train_end(trainer):
+    experiment = comet_ml.get_global_experiment()
+    experiment.log_model("YOLOv8", file_or_folder=trainer.best, file_name="best.pt", overwrite=True)
+
+
+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_train_end": on_train_end} if comet_ml else {}