mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

W&B: Refactor the wandb_utils.py file (#4496)

Browse Source 
* 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

* Fix

* fix

* refactor constructor

* refactor

* refactor

* refactor

* PyCharm reformat

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia 2021-08-23 16:40:07 +05:30 committed by GitHub
parent d1182c4f29
commit 7316b78e36
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 43 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
utils/loggers/wandb/wandb_utils.py
View File

@ -38,6 +38,19 @@ def check_wandb_config_file(data_config_file):
return data_config_file
def check_wandb_dataset(data_file):
is_wandb_artifact = False
if check_file(data_file) and data_file.endswith('.yaml'):
with open(data_file, errors='ignore') as f:
data_dict = yaml.safe_load(f)
is_wandb_artifact = (data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX) or
data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX))
if is_wandb_artifact:
return data_dict
else:
return check_dataset(data_file)
def get_run_info(run_path):
run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
run_id = run_path.stem
@ -147,26 +160,24 @@ class WandbLogger():
allow_val_change=True) if not wandb.run else wandb.run
if self.wandb_run:
if self.job_type == 'Training':
if not opt.resume:
if opt.upload_dataset:
if not opt.resume:
self.wandb_artifact_data_dict = self.check_and_upload_dataset(opt)
elif opt.data.endswith('_wandb.yaml'): # When dataset is W&B artifact
with open(opt.data, errors='ignore') as f:
data_dict = yaml.safe_load(f)
self.data_dict = data_dict
else: # Local .yaml dataset file or .zip file
self.data_dict = check_dataset(opt.data)
if opt.resume:
# resume from artifact
if isinstance(opt.resume, str) and opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
self.data_dict = dict(self.wandb_run.config.data_dict)
else: # local resume
self.data_dict = check_wandb_dataset(opt.data)
else:
self.data_dict = check_dataset(opt.data)
self.data_dict = check_wandb_dataset(opt.data)
self.wandb_artifact_data_dict = self.wandb_artifact_data_dict or self.data_dict
self.setup_training(opt)
if not self.wandb_artifact_data_dict:
self.wandb_artifact_data_dict = self.data_dict
# 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)
self.setup_training(opt)
if self.job_type == 'Dataset Creation':
self.data_dict = self.check_and_upload_dataset(opt)
@ -211,8 +222,6 @@ class WandbLogger():
opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp = str(
self.weights), config.save_period, config.batch_size, config.bbox_interval, config.epochs, \
config.hyp
data_dict = dict(self.wandb_run.config.data_dict) # eliminates the need for config file to resume
else:
data_dict = self.data_dict
if self.val_artifact is None: # If --upload_dataset is set, use the existing artifact, don't download
self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(data_dict.get('train'),