diff --git a/ultralytics/yolo/utils/callbacks/base.py b/ultralytics/yolo/utils/callbacks/base.py index cd12b3c..f0389bf 100644 --- a/ultralytics/yolo/utils/callbacks/base.py +++ b/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) diff --git a/ultralytics/yolo/utils/callbacks/comet.py b/ultralytics/yolo/utils/callbacks/comet.py new file mode 100644 index 0000000..7cb5ce1 --- /dev/null +++ b/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 {}