Comet Logging and Visualization Integration (#9232)
* add comet to logger interface * add comet logger * add support for updated parameters * clean up offline logger creation * update callback args for comet logger * add comet optimizer * add optimizer config * add comet README * update tutorial notebook with Comet section * add option to log class level metrics * add support for class level metrics and confusion matrix * handle errors when adding files to artifacts * fix typo * clean resume workflow * updates for HPO * update comet README * fix typo in comet README * update code snippets in comet README * update comet links in tutorial * updated links * change optimizer batch size param and update comet README image * update comet section in tutorial * use prexisting cmd line flags to configure logger * update artifact upload/download flow * remove come remove comet logger specific cmd line args * move downloading weights into comet logger code * remove extra argparse * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * change checkpoint logging flow to follow offline logger * update resume flow * add comet logger to remote dataset property * update cmd line args in hpo * set types for integer/float env variables * update README * fix typo in README * default to always logging model predictions * Update tutorial.ipynb * Update train.py * Add Comet to Integrations table * Update README.md * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
10 files changed
+1376
-49
No files matched your search
@@ -17,7 +17,7 @@ from utils.loggers.wandb.wandb_utils import WandbLogger
|
||||
from utils.plots import plot_images, plot_labels, plot_results
|
||||
from utils.torch_utils import de_parallel
|
||||
|
||||
LOGGERS = ('csv', 'tb', 'wandb', 'clearml') # *.csv, TensorBoard, Weights & Biases, ClearML
|
||||
LOGGERS = ('csv', 'tb', 'wandb', 'clearml', 'comet') # *.csv, TensorBoard, Weights & Biases, ClearML
|
||||
RANK = int(os.getenv('RANK', -1))
|
||||
|
||||
try:
|
||||
@@ -41,6 +41,18 @@ try:
|
||||
except (ImportError, AssertionError):
|
||||
clearml = None
|
||||
|
||||
try:
|
||||
if RANK not in [0, -1]:
|
||||
comet_ml = None
|
||||
else:
|
||||
import comet_ml
|
||||
|
||||
assert hasattr(comet_ml, '__version__') # verify package import not local dir
|
||||
from utils.loggers.comet import CometLogger
|
||||
|
||||
except (ModuleNotFoundError, ImportError, AssertionError):
|
||||
comet_ml = None
|
||||
|
||||
|
||||
class Loggers():
|
||||
# YOLOv5 Loggers class
|
||||
@@ -80,7 +92,10 @@ class Loggers():
|
||||
prefix = colorstr('ClearML: ')
|
||||
s = f"{prefix}run 'pip install clearml' to automatically track, visualize and remotely train YOLOv5 🚀 in ClearML"
|
||||
self.logger.info(s)
|
||||
|
||||
if not comet_ml:
|
||||
prefix = colorstr('Comet: ')
|
||||
s = f"{prefix}run 'pip install comet_ml' to automatically track and visualize YOLOv5 🚀 runs in Comet"
|
||||
self.logger.info(s)
|
||||
# TensorBoard
|
||||
s = self.save_dir
|
||||
if 'tb' in self.include and not self.opt.evolve:
|
||||
@@ -107,6 +122,18 @@ class Loggers():
|
||||
else:
|
||||
self.clearml = None
|
||||
|
||||
# Comet
|
||||
if comet_ml and 'comet' in self.include:
|
||||
if isinstance(self.opt.resume, str) and self.opt.resume.startswith("comet://"):
|
||||
run_id = self.opt.resume.split("/")[-1]
|
||||
self.comet_logger = CometLogger(self.opt, self.hyp, run_id=run_id)
|
||||
|
||||
else:
|
||||
self.comet_logger = CometLogger(self.opt, self.hyp)
|
||||
|
||||
else:
|
||||
self.comet_logger = None
|
||||
|
||||
@property
|
||||
def remote_dataset(self):
|
||||
# Get data_dict if custom dataset artifact link is provided
|
||||
@@ -115,12 +142,18 @@ class Loggers():
|
||||
data_dict = self.clearml.data_dict
|
||||
if self.wandb:
|
||||
data_dict = self.wandb.data_dict
|
||||
if self.comet_logger:
|
||||
data_dict = self.comet_logger.data_dict
|
||||
|
||||
return data_dict
|
||||
|
||||
def on_train_start(self):
|
||||
# Callback runs on train start
|
||||
pass
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_train_start()
|
||||
|
||||
def on_pretrain_routine_start(self):
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_pretrain_routine_start()
|
||||
|
||||
def on_pretrain_routine_end(self, labels, names):
|
||||
# Callback runs on pre-train routine end
|
||||
@@ -131,8 +164,11 @@ class Loggers():
|
||||
self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})
|
||||
# if self.clearml:
|
||||
# pass # ClearML saves these images automatically using hooks
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_pretrain_routine_end(paths)
|
||||
|
||||
def on_train_batch_end(self, model, ni, imgs, targets, paths):
|
||||
def on_train_batch_end(self, model, ni, imgs, targets, paths, vals):
|
||||
log_dict = dict(zip(self.keys[0:3], vals))
|
||||
# Callback runs on train batch end
|
||||
# ni: number integrated batches (since train start)
|
||||
if self.plots:
|
||||
@@ -148,11 +184,21 @@ class Loggers():
|
||||
if self.clearml:
|
||||
self.clearml.log_debug_samples(files, title='Mosaics')
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_train_batch_end(log_dict, step=ni)
|
||||
|
||||
def on_train_epoch_end(self, epoch):
|
||||
# Callback runs on train epoch end
|
||||
if self.wandb:
|
||||
self.wandb.current_epoch = epoch + 1
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_train_epoch_end(epoch)
|
||||
|
||||
def on_val_start(self):
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_val_start()
|
||||
|
||||
def on_val_image_end(self, pred, predn, path, names, im):
|
||||
# Callback runs on val image end
|
||||
if self.wandb:
|
||||
@@ -160,7 +206,11 @@ class Loggers():
|
||||
if self.clearml:
|
||||
self.clearml.log_image_with_boxes(path, pred, names, im)
|
||||
|
||||
def on_val_end(self):
|
||||
def on_val_batch_end(self, batch_i, im, targets, paths, shapes, out):
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_val_batch_end(batch_i, im, targets, paths, shapes, out)
|
||||
|
||||
def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix):
|
||||
# Callback runs on val end
|
||||
if self.wandb or self.clearml:
|
||||
files = sorted(self.save_dir.glob('val*.jpg'))
|
||||
@@ -169,6 +219,9 @@ class Loggers():
|
||||
if self.clearml:
|
||||
self.clearml.log_debug_samples(files, title='Validation')
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_val_end(nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
|
||||
|
||||
def on_fit_epoch_end(self, vals, epoch, best_fitness, fi):
|
||||
# Callback runs at the end of each fit (train+val) epoch
|
||||
x = dict(zip(self.keys, vals))
|
||||
@@ -199,6 +252,9 @@ class Loggers():
|
||||
self.clearml.current_epoch_logged_images = set() # reset epoch image limit
|
||||
self.clearml.current_epoch += 1
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_fit_epoch_end(x, epoch=epoch)
|
||||
|
||||
def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
|
||||
# Callback runs on model save event
|
||||
if (epoch + 1) % self.opt.save_period == 0 and not final_epoch and self.opt.save_period != -1:
|
||||
@@ -209,6 +265,9 @@ class Loggers():
|
||||
model_name='Latest Model',
|
||||
auto_delete_file=False)
|
||||
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_model_save(last, epoch, final_epoch, best_fitness, fi)
|
||||
|
||||
def on_train_end(self, last, best, epoch, results):
|
||||
# Callback runs on training end, i.e. saving best model
|
||||
if self.plots:
|
||||
@@ -237,10 +296,16 @@ class Loggers():
|
||||
name='Best Model',
|
||||
auto_delete_file=False)
|
||||
|
||||
if self.comet_logger:
|
||||
final_results = dict(zip(self.keys[3:10], results))
|
||||
self.comet_logger.on_train_end(files, self.save_dir, last, best, epoch, final_results)
|
||||
|
||||
def on_params_update(self, params: dict):
|
||||
# Update hyperparams or configs of the experiment
|
||||
if self.wandb:
|
||||
self.wandb.wandb_run.config.update(params, allow_val_change=True)
|
||||
if self.comet_logger:
|
||||
self.comet_logger.on_params_update(params)
|
||||
|
||||
|
||||
class GenericLogger:
|
||||
|
||||
Reference in new issue
Block a user