From 5a70cc2e3f29ea4ecda774031613d27ebc8a0bf5 Mon Sep 17 00:00:00 2001 From: Ayush Chaurasia Date: Mon, 30 Aug 2021 20:37:20 +0530 Subject: [PATCH] W&B: Don't log models in evolve operation (#4611) --- utils/loggers/__init__.py | 12 ++++++++---- utils/loggers/wandb/wandb_utils.py | 2 +- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py index 16d0348d86f3..775803abf068 100644 --- a/utils/loggers/__init__.py +++ b/utils/loggers/__init__.py @@ -138,7 +138,11 @@ def on_train_end(self, last, best, plots, epoch): if self.wandb: self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]}) # Calling wandb.log. TODO: Refactor this into WandbLogger.log_model - wandb.log_artifact(str(best if best.exists() else last), type='model', - name='run_' + self.wandb.wandb_run.id + '_model', - aliases=['latest', 'best', 'stripped']) - self.wandb.finish_run() + if not self.opt.evolve: + wandb.log_artifact(str(best if best.exists() else last), type='model', + name='run_' + self.wandb.wandb_run.id + '_model', + aliases=['latest', 'best', 'stripped']) + self.wandb.finish_run() + else: + self.wandb.finish_run() + self.wandb = WandbLogger(self.opt) diff --git a/utils/loggers/wandb/wandb_utils.py b/utils/loggers/wandb/wandb_utils.py index 8b2095afcb8b..5d495c70517b 100644 --- a/utils/loggers/wandb/wandb_utils.py +++ b/utils/loggers/wandb/wandb_utils.py @@ -112,7 +112,7 @@ class WandbLogger(): https://docs.wandb.com/guides/integrations/yolov5 """ - def __init__(self, opt, run_id, job_type='Training'): + def __init__(self, opt, run_id=None, job_type='Training'): """ - Initialize WandbLogger instance - Upload dataset if opt.upload_dataset is True