Unverified 提交 0cae7576 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

utils/wandb_logging PEP8 reformat (#2755)

* wandb_logging PEP8 reformat * Update wandb_utils.py
上级 b5de52c4
import argparse import argparse
from pathlib import Path
import yaml import yaml
from wandb_utils import WandbLogger from wandb_utils import WandbLogger
from utils.datasets import LoadImagesAndLabels
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://' WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
...@@ -21,6 +19,6 @@ if __name__ == '__main__': ...@@ -21,6 +19,6 @@ if __name__ == '__main__':
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') 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') parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project')
opt = parser.parse_args() 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) create_dataset_artifact(opt)
import argparse
import json import json
import os
import shutil
import sys import sys
from pathlib import Path
import torch import torch
import yaml import yaml
from datetime import datetime
from pathlib import Path
from tqdm import tqdm from tqdm import tqdm
sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path
...@@ -33,6 +30,7 @@ def check_wandb_config_file(data_config_file): ...@@ -33,6 +30,7 @@ def check_wandb_config_file(data_config_file):
return wandb_config return wandb_config
return data_config_file return data_config_file
def get_run_info(run_path): def get_run_info(run_path):
run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX)) run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
run_id = run_path.stem run_id = run_path.stem
...@@ -40,11 +38,12 @@ def get_run_info(run_path): ...@@ -40,11 +38,12 @@ def get_run_info(run_path):
model_artifact_name = 'run_' + run_id + '_model' model_artifact_name = 'run_' + run_id + '_model'
return run_id, project, model_artifact_name return run_id, project, model_artifact_name
def check_wandb_resume(opt): def check_wandb_resume(opt):
process_wandb_config_ddp_mode(opt) if opt.global_rank not in [-1, 0] else None process_wandb_config_ddp_mode(opt) if opt.global_rank not in [-1, 0] else None
if isinstance(opt.resume, str): if isinstance(opt.resume, str):
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX): 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) run_id, project, model_artifact_name = get_run_info(opt.resume)
api = wandb.Api() api = wandb.Api()
artifact = api.artifact(project + '/' + model_artifact_name + ':latest') artifact = api.artifact(project + '/' + model_artifact_name + ':latest')
...@@ -53,6 +52,7 @@ def check_wandb_resume(opt): ...@@ -53,6 +52,7 @@ def check_wandb_resume(opt):
return True return True
return None return None
def process_wandb_config_ddp_mode(opt): def process_wandb_config_ddp_mode(opt):
with open(opt.data) as f: with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
...@@ -63,7 +63,7 @@ def process_wandb_config_ddp_mode(opt): ...@@ -63,7 +63,7 @@ def process_wandb_config_ddp_mode(opt):
train_dir = train_artifact.download() train_dir = train_artifact.download()
train_path = Path(train_dir) / 'data/images/' train_path = Path(train_dir) / 'data/images/'
data_dict['train'] = str(train_path) data_dict['train'] = str(train_path)
if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX): if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX):
api = wandb.Api() api = wandb.Api()
val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias) val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias)
...@@ -71,12 +71,11 @@ def process_wandb_config_ddp_mode(opt): ...@@ -71,12 +71,11 @@ def process_wandb_config_ddp_mode(opt):
val_path = Path(val_dir) / 'data/images/' val_path = Path(val_dir) / 'data/images/'
data_dict['val'] = str(val_path) data_dict['val'] = str(val_path)
if train_dir or val_dir: 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: with open(ddp_data_path, 'w') as f:
yaml.dump(data_dict, f) yaml.dump(data_dict, f)
opt.data = ddp_data_path opt.data = ddp_data_path
class WandbLogger(): class WandbLogger():
def __init__(self, opt, name, run_id, data_dict, job_type='Training'): def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
...@@ -84,7 +83,7 @@ class WandbLogger(): ...@@ -84,7 +83,7 @@ class WandbLogger():
self.job_type = job_type self.job_type = job_type
self.wandb, self.wandb_run, self.data_dict = wandb, None if not wandb else wandb.run, data_dict 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 # 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): if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
run_id, project, model_artifact_name = get_run_info(opt.resume) run_id, project, model_artifact_name = get_run_info(opt.resume)
model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
...@@ -98,7 +97,7 @@ class WandbLogger(): ...@@ -98,7 +97,7 @@ class WandbLogger():
project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem, project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
name=name, name=name,
job_type=job_type, 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.wandb_run:
if self.job_type == 'Training': if self.job_type == 'Training':
if not opt.resume: if not opt.resume:
...@@ -110,15 +109,15 @@ class WandbLogger(): ...@@ -110,15 +109,15 @@ class WandbLogger():
if self.job_type == 'Dataset Creation': if self.job_type == 'Dataset Creation':
self.data_dict = self.check_and_upload_dataset(opt) self.data_dict = self.check_and_upload_dataset(opt)
else: 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): def check_and_upload_dataset(self, opt):
assert wandb, 'Install wandb to upload dataset' assert wandb, 'Install wandb to upload dataset'
check_dataset(self.data_dict) check_dataset(self.data_dict)
config_path = self.log_dataset_artifact(opt.data, config_path = self.log_dataset_artifact(opt.data,
opt.single_cls, opt.single_cls,
'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem) 'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem)
print("Created dataset config file ", config_path) print("Created dataset config file ", config_path)
with open(config_path) as f: with open(config_path) as f:
wandb_data_dict = yaml.load(f, Loader=yaml.SafeLoader) wandb_data_dict = yaml.load(f, Loader=yaml.SafeLoader)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论