Skip to content

Commit

Permalink
utils/wandb_logging PEP8 reformat (ultralytics#2755)
Browse files Browse the repository at this point in the history
* wandb_logging PEP8 reformat

* Update wandb_utils.py
  • Loading branch information
glenn-jocher committed Apr 10, 2021
1 parent 95b6334 commit 079d98c
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 20 deletions.
6 changes: 2 additions & 4 deletions utils/wandb_logging/log_dataset.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,8 @@
import argparse
from pathlib import Path

import yaml

from wandb_utils import WandbLogger
from utils.datasets import LoadImagesAndLabels

WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'

Expand All @@ -21,6 +19,6 @@ def create_dataset_artifact(opt):
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project')
opt = parser.parse_args()
opt.resume = False # Explicitly disallow resume check for dataset upload Job
opt.resume = False # Explicitly disallow resume check for dataset upload job

create_dataset_artifact(opt)
31 changes: 15 additions & 16 deletions utils/wandb_logging/wandb_utils.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,9 @@
import argparse
import json
import os
import shutil
import sys
from pathlib import Path

import torch
import yaml
from datetime import datetime
from pathlib import Path
from tqdm import tqdm

sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path
Expand All @@ -33,18 +30,20 @@ def check_wandb_config_file(data_config_file):
return wandb_config
return data_config_file


def get_run_info(run_path):
run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
run_id = run_path.stem
project = run_path.parent.stem
model_artifact_name = 'run_' + run_id + '_model'
return run_id, project, model_artifact_name


def check_wandb_resume(opt):
process_wandb_config_ddp_mode(opt) if opt.global_rank not in [-1, 0] else None
if isinstance(opt.resume, str):
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
if opt.global_rank not in [-1, 0]: # For resuming DDP runs
if opt.global_rank not in [-1, 0]: # For resuming DDP runs
run_id, project, model_artifact_name = get_run_info(opt.resume)
api = wandb.Api()
artifact = api.artifact(project + '/' + model_artifact_name + ':latest')
Expand All @@ -53,6 +52,7 @@ def check_wandb_resume(opt):
return True
return None


def process_wandb_config_ddp_mode(opt):
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
Expand All @@ -63,28 +63,27 @@ def process_wandb_config_ddp_mode(opt):
train_dir = train_artifact.download()
train_path = Path(train_dir) / 'data/images/'
data_dict['train'] = str(train_path)

if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX):
api = wandb.Api()
val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias)
val_dir = val_artifact.download()
val_path = Path(val_dir) / 'data/images/'
data_dict['val'] = str(val_path)
if train_dir or val_dir:
ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml')
ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml')
with open(ddp_data_path, 'w') as f:
yaml.dump(data_dict, f)
opt.data = ddp_data_path




class WandbLogger():
def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
# Pre-training routine --
self.job_type = job_type
self.wandb, self.wandb_run, self.data_dict = wandb, None if not wandb else wandb.run, data_dict
# It's more elegant to stick to 1 wandb.init call, but useful config data is overwritten in the WandbLogger's wandb.init call
if isinstance(opt.resume, str): # checks resume from artifact
if isinstance(opt.resume, str): # checks resume from artifact
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
run_id, project, model_artifact_name = get_run_info(opt.resume)
model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
Expand All @@ -98,7 +97,7 @@ def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
name=name,
job_type=job_type,
id=run_id) if not wandb.run else wandb.run
id=run_id) if not wandb.run else wandb.run
if self.wandb_run:
if self.job_type == 'Training':
if not opt.resume:
Expand All @@ -110,15 +109,15 @@ def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
if self.job_type == 'Dataset Creation':
self.data_dict = self.check_and_upload_dataset(opt)
else:
print(f"{colorstr('wandb: ')}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")

prefix = colorstr('wandb: ')
print(f"{prefix}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")

def check_and_upload_dataset(self, opt):
assert wandb, 'Install wandb to upload dataset'
check_dataset(self.data_dict)
config_path = self.log_dataset_artifact(opt.data,
opt.single_cls,
'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem)
opt.single_cls,
'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem)
print("Created dataset config file ", config_path)
with open(config_path) as f:
wandb_data_dict = yaml.load(f, Loader=yaml.SafeLoader)
Expand Down

0 comments on commit 079d98c

Please sign in to comment.