W&B: More improvements and refactoring (ultralytics#4205)

* Improve docstrings and run names * default wandb login prompt with timeout * return key * Update api_key check logic * Properly support zipped dataset feature * update docstring * Revert tuorial change * extend changes to log_dataset * add run name * bug fix * bug fix * Update comment * fix import check * remove unused import * Hardcore .yaml file extension * reduce code * Reformat using pycharm * Remove redundant try catch * More refactoring and bug fixes * retry * Reformat using pycharm * respect LOGGERS include list Co-authored-by: Glenn Jocher <[email protected]>
fhkiel-mlaip · Jul 28, 2021 · 6a54bb7 · 6a54bb7
1 parent d20ee91
commit 6a54bb7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py
@@ -48,12 +48,12 @@ def start(self):
             self.tb = SummaryWriter(str(s))
 
         # W&B
-        try:
-            assert 'wandb' in self.include and wandb
-            run_id = torch.load(self.weights).get('wandb_id') if self.opt.resume else None
+        if wandb and 'wandb' in self.include:
+            wandb_artifact_resume = isinstance(self.opt.resume, str) and self.opt.resume.startswith('wandb-artifact://')
+            run_id = torch.load(self.weights).get('wandb_id') if self.opt.resume and not wandb_artifact_resume else None
             self.opt.hyp = self.hyp  # add hyperparameters
             self.wandb = WandbLogger(self.opt, run_id)
-        except:
+        else:
             self.wandb = None
 
         return self

diff --git a/utils/loggers/wandb/wandb_utils.py b/utils/loggers/wandb/wandb_utils.py
@@ -158,11 +158,12 @@ def __init__(self, opt, run_id, job_type='Training'):
                         self.data_dict = check_dataset(opt.data)
 
                 self.setup_training(opt)
-                # write data_dict to config. useful for resuming from artifacts
                 if not self.wandb_artifact_data_dict:
                     self.wandb_artifact_data_dict = self.data_dict
-                self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict},
-                                             allow_val_change=True)
+                # write data_dict to config. useful for resuming from artifacts. Do this only when not resuming.
+                if not opt.resume:
+                    self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict},
+                                                 allow_val_change=True)
 
             if self.job_type == 'Dataset Creation':
                 self.data_dict = self.check_and_upload_dataset(opt)