W&B: Don't log models in evolve operation (#4611)

ultralytics · Aug 30, 2021 · 11e7c7b · 11e7c7b
1 parent 3a72d4a
commit 11e7c7b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git 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
@@ -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