utils/wandb_logging PEP8 reformat (#2755)
* wandb_logging PEP8 reformat * Update wandb_utils.py
This commit is contained in:
parent
b5de52c4cd
commit
0cae7576a9
|
|
@ -1,10 +1,8 @@
|
||||||
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__':
|
||||||
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)
|
||||||
|
|
|
||||||
|
|
@ -1,12 +1,9 @@
|
||||||
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):
|
||||||
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):
|
||||||
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):
|
||||||
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):
|
||||||
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):
|
||||||
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():
|
||||||
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():
|
||||||
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():
|
||||||
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)
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue