Refactor train.py and val.py loggers (#4137)
* Update loggers * Config * Update val.py * cleanup * fix1 * fix2 * fix3 and reformat * format sweep.py * Logger() class * cleanup * cleanup2 * wandb package import fix * wandb package import fix2 * txt fix * fix4 * fix5 * fix6 * drop wandb into utils/loggers * fix 7 * rename loggers/wandb_logging to loggers/wandb * Update message * Update message * Update message * cleanup * Fix x axis bug * fix rank 0 issue * cleanup
This commit is contained in:
Glenn Jocher authored and GitHub committed
1 parent
63dd65e7ed
commit
efe60b5681
9 files changed
+172
-91
No files matched your search
@@ -0,0 +1,129 @@
|
||||
# YOLOv5 experiment logging utils
|
||||
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from utils.general import colorstr, emojis
|
||||
from utils.loggers.wandb.wandb_utils import WandbLogger
|
||||
from utils.torch_utils import de_parallel
|
||||
|
||||
LOGGERS = ('txt', 'tb', 'wandb') # text-file, TensorBoard, Weights & Biases
|
||||
|
||||
try:
|
||||
import wandb
|
||||
|
||||
assert hasattr(wandb, '__version__') # verify package import not local dir
|
||||
except (ImportError, AssertionError):
|
||||
wandb = None
|
||||
|
||||
|
||||
class Loggers():
|
||||
# YOLOv5 Loggers class
|
||||
def __init__(self, save_dir=None, results_file=None, weights=None, opt=None, hyp=None,
|
||||
data_dict=None, logger=None, include=LOGGERS):
|
||||
self.save_dir = save_dir
|
||||
self.results_file = results_file
|
||||
self.weights = weights
|
||||
self.opt = opt
|
||||
self.hyp = hyp
|
||||
self.data_dict = data_dict
|
||||
self.logger = logger # for printing results to console
|
||||
self.include = include
|
||||
for k in LOGGERS:
|
||||
setattr(self, k, None) # init empty logger dictionary
|
||||
|
||||
def start(self):
|
||||
self.txt = True # always log to txt
|
||||
|
||||
# Message
|
||||
try:
|
||||
import wandb
|
||||
except ImportError:
|
||||
prefix = colorstr('Weights & Biases: ')
|
||||
s = f"{prefix}run 'pip install wandb' to automatically track and visualize YOLOv5 🚀 runs (RECOMMENDED)"
|
||||
print(emojis(s))
|
||||
|
||||
# TensorBoard
|
||||
s = self.save_dir
|
||||
if 'tb' in self.include and not self.opt.evolve:
|
||||
prefix = colorstr('TensorBoard: ')
|
||||
self.logger.info(f"{prefix}Start with 'tensorboard --logdir {s.parent}', view at http://localhost:6006/")
|
||||
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
|
||||
self.opt.hyp = self.hyp # add hyperparameters
|
||||
self.wandb = WandbLogger(self.opt, s.stem, run_id, self.data_dict)
|
||||
except:
|
||||
self.wandb = None
|
||||
|
||||
return self
|
||||
|
||||
def on_train_batch_end(self, ni, model, imgs):
|
||||
# Callback runs on train batch end
|
||||
if ni == 0:
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter('ignore') # suppress jit trace warning
|
||||
self.tb.add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), [])
|
||||
if self.wandb and ni == 10:
|
||||
files = sorted(self.save_dir.glob('train*.jpg'))
|
||||
self.wandb.log({'Mosaics': [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]})
|
||||
|
||||
def on_train_epoch_end(self, epoch):
|
||||
# Callback runs on train epoch end
|
||||
if self.wandb:
|
||||
self.wandb.current_epoch = epoch + 1
|
||||
|
||||
def on_val_batch_end(self, pred, predn, path, names, im):
|
||||
# Callback runs on train batch end
|
||||
if self.wandb:
|
||||
self.wandb.val_one_image(pred, predn, path, names, im)
|
||||
|
||||
def on_val_end(self):
|
||||
# Callback runs on val end
|
||||
if self.wandb:
|
||||
files = sorted(self.save_dir.glob('val*.jpg'))
|
||||
self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]})
|
||||
|
||||
def on_train_val_end(self, mloss, results, lr, epoch, s, best_fitness, fi):
|
||||
# Callback runs on validation end during training
|
||||
vals = list(mloss[:-1]) + list(results) + lr
|
||||
tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
|
||||
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
|
||||
'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss
|
||||
'x/lr0', 'x/lr1', 'x/lr2'] # params
|
||||
if self.txt:
|
||||
with open(self.results_file, 'a') as f:
|
||||
f.write(s + '%10.4g' * 7 % results + '\n') # append metrics, val_loss
|
||||
if self.tb:
|
||||
for x, tag in zip(vals, tags):
|
||||
self.tb.add_scalar(tag, x, epoch) # TensorBoard
|
||||
if self.wandb:
|
||||
self.wandb.log({k: v for k, v in zip(tags, vals)})
|
||||
self.wandb.end_epoch(best_result=best_fitness == fi)
|
||||
|
||||
def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
|
||||
# Callback runs on model save event
|
||||
if self.wandb:
|
||||
if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
|
||||
self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
|
||||
|
||||
def on_train_end(self, last, best):
|
||||
# Callback runs on training end
|
||||
files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
|
||||
files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()] # filter
|
||||
if self.wandb:
|
||||
wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
|
||||
wandb.log_artifact(str(best if best.exists() else last), type='model',
|
||||
name='run_' + self.wandb.wandb_run.id + '_model',
|
||||
aliases=['latest', 'best', 'stripped'])
|
||||
self.wandb.finish_run()
|
||||
|
||||
def log_images(self, paths):
|
||||
# Log images
|
||||
if self.wandb:
|
||||
self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})
|
||||
Whitespace-only changes.
@@ -0,0 +1,26 @@
|
||||
import argparse
|
||||
|
||||
import yaml
|
||||
|
||||
from wandb_utils import WandbLogger
|
||||
|
||||
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
|
||||
|
||||
|
||||
def create_dataset_artifact(opt):
|
||||
with open(opt.data) as f:
|
||||
data = yaml.safe_load(f) # data dict
|
||||
logger = WandbLogger(opt, '', None, data, job_type='Dataset Creation')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
|
||||
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('--entity', default=None, help='W&B entity')
|
||||
|
||||
opt = parser.parse_args()
|
||||
opt.resume = False # Explicitly disallow resume check for dataset upload job
|
||||
|
||||
create_dataset_artifact(opt)
|
||||
@@ -0,0 +1,33 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import wandb
|
||||
|
||||
FILE = Path(__file__).absolute()
|
||||
sys.path.append(FILE.parents[2].as_posix()) # add utils/ to path
|
||||
|
||||
from train import train, parse_opt
|
||||
from utils.general import increment_path
|
||||
from utils.torch_utils import select_device
|
||||
|
||||
|
||||
def sweep():
|
||||
wandb.init()
|
||||
# Get hyp dict from sweep agent
|
||||
hyp_dict = vars(wandb.config).get("_items")
|
||||
|
||||
# Workaround: get necessary opt args
|
||||
opt = parse_opt(known=True)
|
||||
opt.batch_size = hyp_dict.get("batch_size")
|
||||
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok or opt.evolve))
|
||||
opt.epochs = hyp_dict.get("epochs")
|
||||
opt.nosave = True
|
||||
opt.data = hyp_dict.get("data")
|
||||
device = select_device(opt.device, batch_size=opt.batch_size)
|
||||
|
||||
# train
|
||||
train(hyp_dict, opt, device)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sweep()
|
||||
@@ -0,0 +1,143 @@
|
||||
# Hyperparameters for training
|
||||
# To set range-
|
||||
# Provide min and max values as:
|
||||
# parameter:
|
||||
#
|
||||
# min: scalar
|
||||
# max: scalar
|
||||
# OR
|
||||
#
|
||||
# Set a specific list of search space-
|
||||
# parameter:
|
||||
# values: [scalar1, scalar2, scalar3...]
|
||||
#
|
||||
# You can use grid, bayesian and hyperopt search strategy
|
||||
# For more info on configuring sweeps visit - https://docs.wandb.ai/guides/sweeps/configuration
|
||||
|
||||
program: utils/loggers/wandb/sweep.py
|
||||
method: random
|
||||
metric:
|
||||
name: metrics/mAP_0.5
|
||||
goal: maximize
|
||||
|
||||
parameters:
|
||||
# hyperparameters: set either min, max range or values list
|
||||
data:
|
||||
value: "data/coco128.yaml"
|
||||
batch_size:
|
||||
values: [ 64 ]
|
||||
epochs:
|
||||
values: [ 10 ]
|
||||
|
||||
lr0:
|
||||
distribution: uniform
|
||||
min: 1e-5
|
||||
max: 1e-1
|
||||
lrf:
|
||||
distribution: uniform
|
||||
min: 0.01
|
||||
max: 1.0
|
||||
momentum:
|
||||
distribution: uniform
|
||||
min: 0.6
|
||||
max: 0.98
|
||||
weight_decay:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.001
|
||||
warmup_epochs:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 5.0
|
||||
warmup_momentum:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.95
|
||||
warmup_bias_lr:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.2
|
||||
box:
|
||||
distribution: uniform
|
||||
min: 0.02
|
||||
max: 0.2
|
||||
cls:
|
||||
distribution: uniform
|
||||
min: 0.2
|
||||
max: 4.0
|
||||
cls_pw:
|
||||
distribution: uniform
|
||||
min: 0.5
|
||||
max: 2.0
|
||||
obj:
|
||||
distribution: uniform
|
||||
min: 0.2
|
||||
max: 4.0
|
||||
obj_pw:
|
||||
distribution: uniform
|
||||
min: 0.5
|
||||
max: 2.0
|
||||
iou_t:
|
||||
distribution: uniform
|
||||
min: 0.1
|
||||
max: 0.7
|
||||
anchor_t:
|
||||
distribution: uniform
|
||||
min: 2.0
|
||||
max: 8.0
|
||||
fl_gamma:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.1
|
||||
hsv_h:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.1
|
||||
hsv_s:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.9
|
||||
hsv_v:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.9
|
||||
degrees:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 45.0
|
||||
translate:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.9
|
||||
scale:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.9
|
||||
shear:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 10.0
|
||||
perspective:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 0.001
|
||||
flipud:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 1.0
|
||||
fliplr:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 1.0
|
||||
mosaic:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 1.0
|
||||
mixup:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 1.0
|
||||
copy_paste:
|
||||
distribution: uniform
|
||||
min: 0.0
|
||||
max: 1.0
|
||||
@@ -0,0 +1,374 @@
|
||||
"""Utilities and tools for tracking runs with Weights & Biases."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from tqdm import tqdm
|
||||
|
||||
FILE = Path(__file__).absolute()
|
||||
sys.path.append(FILE.parents[3].as_posix()) # add yolov5/ to path
|
||||
|
||||
from utils.datasets import LoadImagesAndLabels
|
||||
from utils.datasets import img2label_paths
|
||||
from utils.general import check_dataset, check_file
|
||||
|
||||
try:
|
||||
import wandb
|
||||
|
||||
assert hasattr(wandb, '__version__') # verify package import not local dir
|
||||
except (ImportError, AssertionError):
|
||||
wandb = None
|
||||
|
||||
RANK = int(os.getenv('RANK', -1))
|
||||
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
|
||||
|
||||
|
||||
def remove_prefix(from_string, prefix=WANDB_ARTIFACT_PREFIX):
|
||||
return from_string[len(prefix):]
|
||||
|
||||
|
||||
def check_wandb_config_file(data_config_file):
|
||||
wandb_config = '_wandb.'.join(data_config_file.rsplit('.', 1)) # updated data.yaml path
|
||||
if Path(wandb_config).is_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
|
||||
entity = run_path.parent.parent.stem
|
||||
model_artifact_name = 'run_' + run_id + '_model'
|
||||
return entity, project, run_id, model_artifact_name
|
||||
|
||||
|
||||
def check_wandb_resume(opt):
|
||||
process_wandb_config_ddp_mode(opt) if RANK not in [-1, 0] else None
|
||||
if isinstance(opt.resume, str):
|
||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
if RANK not in [-1, 0]: # For resuming DDP runs
|
||||
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
||||
api = wandb.Api()
|
||||
artifact = api.artifact(entity + '/' + project + '/' + model_artifact_name + ':latest')
|
||||
modeldir = artifact.download()
|
||||
opt.weights = str(Path(modeldir) / "last.pt")
|
||||
return True
|
||||
return None
|
||||
|
||||
|
||||
def process_wandb_config_ddp_mode(opt):
|
||||
with open(check_file(opt.data)) as f:
|
||||
data_dict = yaml.safe_load(f) # data dict
|
||||
train_dir, val_dir = None, None
|
||||
if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX):
|
||||
api = wandb.Api()
|
||||
train_artifact = api.artifact(remove_prefix(data_dict['train']) + ':' + opt.artifact_alias)
|
||||
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')
|
||||
with open(ddp_data_path, 'w') as f:
|
||||
yaml.safe_dump(data_dict, f)
|
||||
opt.data = ddp_data_path
|
||||
|
||||
|
||||
class WandbLogger():
|
||||
"""Log training runs, datasets, models, and predictions to Weights & Biases.
|
||||
|
||||
This logger sends information to W&B at wandb.ai. By default, this information
|
||||
includes hyperparameters, system configuration and metrics, model metrics,
|
||||
and basic data metrics and analyses.
|
||||
|
||||
By providing additional command line arguments to train.py, datasets,
|
||||
models and predictions can also be logged.
|
||||
|
||||
For more on how this logger is used, see the Weights & Biases documentation:
|
||||
https://docs.wandb.com/guides/integrations/yolov5
|
||||
"""
|
||||
|
||||
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 = wandb, None if not wandb else wandb.run
|
||||
self.val_artifact, self.train_artifact = None, None
|
||||
self.train_artifact_path, self.val_artifact_path = None, None
|
||||
self.result_artifact = None
|
||||
self.val_table, self.result_table = None, None
|
||||
self.data_dict = data_dict
|
||||
self.bbox_media_panel_images = []
|
||||
self.val_table_path_map = None
|
||||
self.max_imgs_to_log = 16
|
||||
# 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 opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
||||
model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
|
||||
assert wandb, 'install wandb to resume wandb runs'
|
||||
# Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config
|
||||
self.wandb_run = wandb.init(id=run_id,
|
||||
project=project,
|
||||
entity=entity,
|
||||
resume='allow',
|
||||
allow_val_change=True)
|
||||
opt.resume = model_artifact_name
|
||||
elif self.wandb:
|
||||
self.wandb_run = wandb.init(config=opt,
|
||||
resume="allow",
|
||||
project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
|
||||
entity=opt.entity,
|
||||
name=name,
|
||||
job_type=job_type,
|
||||
id=run_id,
|
||||
allow_val_change=True) if not wandb.run else wandb.run
|
||||
if self.wandb_run:
|
||||
if self.job_type == 'Training':
|
||||
if not opt.resume:
|
||||
wandb_data_dict = self.check_and_upload_dataset(opt) if opt.upload_dataset else data_dict
|
||||
# Info useful for resuming from artifacts
|
||||
self.wandb_run.config.update({'opt': vars(opt), 'data_dict': wandb_data_dict},
|
||||
allow_val_change=True)
|
||||
self.data_dict = self.setup_training(opt, data_dict)
|
||||
if self.job_type == 'Dataset Creation':
|
||||
self.data_dict = self.check_and_upload_dataset(opt)
|
||||
|
||||
def check_and_upload_dataset(self, opt):
|
||||
assert wandb, 'Install wandb to upload dataset'
|
||||
config_path = self.log_dataset_artifact(check_file(opt.data),
|
||||
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.safe_load(f)
|
||||
return wandb_data_dict
|
||||
|
||||
def setup_training(self, opt, data_dict):
|
||||
self.log_dict, self.current_epoch = {}, 0
|
||||
self.bbox_interval = opt.bbox_interval
|
||||
if isinstance(opt.resume, str):
|
||||
modeldir, _ = self.download_model_artifact(opt)
|
||||
if modeldir:
|
||||
self.weights = Path(modeldir) / "last.pt"
|
||||
config = self.wandb_run.config
|
||||
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.opt['hyp']
|
||||
data_dict = dict(self.wandb_run.config.data_dict) # eliminates the need for config file to resume
|
||||
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'),
|
||||
opt.artifact_alias)
|
||||
self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'),
|
||||
opt.artifact_alias)
|
||||
|
||||
if self.train_artifact_path is not None:
|
||||
train_path = Path(self.train_artifact_path) / 'data/images/'
|
||||
data_dict['train'] = str(train_path)
|
||||
if self.val_artifact_path is not None:
|
||||
val_path = Path(self.val_artifact_path) / 'data/images/'
|
||||
data_dict['val'] = str(val_path)
|
||||
|
||||
if self.val_artifact is not None:
|
||||
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
||||
self.result_table = wandb.Table(["epoch", "id", "ground truth", "prediction", "avg_confidence"])
|
||||
self.val_table = self.val_artifact.get("val")
|
||||
if self.val_table_path_map is None:
|
||||
self.map_val_table_path()
|
||||
wandb.log({"validation dataset": self.val_table})
|
||||
if opt.bbox_interval == -1:
|
||||
self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1
|
||||
return data_dict
|
||||
|
||||
def download_dataset_artifact(self, path, alias):
|
||||
if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
artifact_path = Path(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
|
||||
dataset_artifact = wandb.use_artifact(artifact_path.as_posix().replace("\\", "/"))
|
||||
assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
|
||||
datadir = dataset_artifact.download()
|
||||
return datadir, dataset_artifact
|
||||
return None, None
|
||||
|
||||
def download_model_artifact(self, opt):
|
||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
model_artifact = wandb.use_artifact(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX) + ":latest")
|
||||
assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
|
||||
modeldir = model_artifact.download()
|
||||
epochs_trained = model_artifact.metadata.get('epochs_trained')
|
||||
total_epochs = model_artifact.metadata.get('total_epochs')
|
||||
is_finished = total_epochs is None
|
||||
assert not is_finished, 'training is finished, can only resume incomplete runs.'
|
||||
return modeldir, model_artifact
|
||||
return None, None
|
||||
|
||||
def log_model(self, path, opt, epoch, fitness_score, best_model=False):
|
||||
model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
|
||||
'original_url': str(path),
|
||||
'epochs_trained': epoch + 1,
|
||||
'save period': opt.save_period,
|
||||
'project': opt.project,
|
||||
'total_epochs': opt.epochs,
|
||||
'fitness_score': fitness_score
|
||||
})
|
||||
model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
|
||||
wandb.log_artifact(model_artifact,
|
||||
aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
|
||||
print("Saving model artifact on epoch ", epoch + 1)
|
||||
|
||||
def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False):
|
||||
with open(data_file) as f:
|
||||
data = yaml.safe_load(f) # data dict
|
||||
check_dataset(data)
|
||||
nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names'])
|
||||
names = {k: v for k, v in enumerate(names)} # to index dictionary
|
||||
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
||||
data['train'], rect=True, batch_size=1), names, name='train') if data.get('train') else None
|
||||
self.val_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
||||
data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None
|
||||
if data.get('train'):
|
||||
data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train')
|
||||
if data.get('val'):
|
||||
data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val')
|
||||
path = data_file if overwrite_config else '_wandb.'.join(data_file.rsplit('.', 1)) # updated data.yaml path
|
||||
data.pop('download', None)
|
||||
data.pop('path', None)
|
||||
with open(path, 'w') as f:
|
||||
yaml.safe_dump(data, f)
|
||||
|
||||
if self.job_type == 'Training': # builds correct artifact pipeline graph
|
||||
self.wandb_run.use_artifact(self.val_artifact)
|
||||
self.wandb_run.use_artifact(self.train_artifact)
|
||||
self.val_artifact.wait()
|
||||
self.val_table = self.val_artifact.get('val')
|
||||
self.map_val_table_path()
|
||||
else:
|
||||
self.wandb_run.log_artifact(self.train_artifact)
|
||||
self.wandb_run.log_artifact(self.val_artifact)
|
||||
return path
|
||||
|
||||
def map_val_table_path(self):
|
||||
self.val_table_path_map = {}
|
||||
print("Mapping dataset")
|
||||
for i, data in enumerate(tqdm(self.val_table.data)):
|
||||
self.val_table_path_map[data[3]] = data[0]
|
||||
|
||||
def create_dataset_table(self, dataset, class_to_id, name='dataset'):
|
||||
# TODO: Explore multiprocessing to slpit this loop parallely| This is essential for speeding up the the logging
|
||||
artifact = wandb.Artifact(name=name, type="dataset")
|
||||
img_files = tqdm([dataset.path]) if isinstance(dataset.path, str) and Path(dataset.path).is_dir() else None
|
||||
img_files = tqdm(dataset.img_files) if not img_files else img_files
|
||||
for img_file in img_files:
|
||||
if Path(img_file).is_dir():
|
||||
artifact.add_dir(img_file, name='data/images')
|
||||
labels_path = 'labels'.join(dataset.path.rsplit('images', 1))
|
||||
artifact.add_dir(labels_path, name='data/labels')
|
||||
else:
|
||||
artifact.add_file(img_file, name='data/images/' + Path(img_file).name)
|
||||
label_file = Path(img2label_paths([img_file])[0])
|
||||
artifact.add_file(str(label_file),
|
||||
name='data/labels/' + label_file.name) if label_file.exists() else None
|
||||
table = wandb.Table(columns=["id", "train_image", "Classes", "name"])
|
||||
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
|
||||
for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
|
||||
box_data, img_classes = [], {}
|
||||
for cls, *xywh in labels[:, 1:].tolist():
|
||||
cls = int(cls)
|
||||
box_data.append({"position": {"middle": [xywh[0], xywh[1]], "width": xywh[2], "height": xywh[3]},
|
||||
"class_id": cls,
|
||||
"box_caption": "%s" % (class_to_id[cls])})
|
||||
img_classes[cls] = class_to_id[cls]
|
||||
boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
|
||||
table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()),
|
||||
Path(paths).name)
|
||||
artifact.add(table, name)
|
||||
return artifact
|
||||
|
||||
def log_training_progress(self, predn, path, names):
|
||||
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in names.items()])
|
||||
box_data = []
|
||||
total_conf = 0
|
||||
for *xyxy, conf, cls in predn.tolist():
|
||||
if conf >= 0.25:
|
||||
box_data.append(
|
||||
{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||
"class_id": int(cls),
|
||||
"box_caption": "%s %.3f" % (names[cls], conf),
|
||||
"scores": {"class_score": conf},
|
||||
"domain": "pixel"})
|
||||
total_conf = total_conf + conf
|
||||
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
||||
id = self.val_table_path_map[Path(path).name]
|
||||
self.result_table.add_data(self.current_epoch,
|
||||
id,
|
||||
self.val_table.data[id][1],
|
||||
wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set),
|
||||
total_conf / max(1, len(box_data))
|
||||
)
|
||||
|
||||
def val_one_image(self, pred, predn, path, names, im):
|
||||
if self.val_table and self.result_table: # Log Table if Val dataset is uploaded as artifact
|
||||
self.log_training_progress(predn, path, names)
|
||||
else: # Default to bbox media panelif Val artifact not found
|
||||
if len(self.bbox_media_panel_images) < self.max_imgs_to_log and self.current_epoch > 0:
|
||||
if self.current_epoch % self.bbox_interval == 0:
|
||||
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||
"class_id": int(cls),
|
||||
"box_caption": "%s %.3f" % (names[cls], conf),
|
||||
"scores": {"class_score": conf},
|
||||
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
|
||||
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
||||
self.bbox_media_panel_images.append(wandb.Image(im, boxes=boxes, caption=path.name))
|
||||
|
||||
def log(self, log_dict):
|
||||
if self.wandb_run:
|
||||
for key, value in log_dict.items():
|
||||
self.log_dict[key] = value
|
||||
|
||||
def end_epoch(self, best_result=False):
|
||||
if self.wandb_run:
|
||||
with all_logging_disabled():
|
||||
if self.bbox_media_panel_images:
|
||||
self.log_dict["Bounding Box Debugger/Images"] = self.bbox_media_panel_images
|
||||
wandb.log(self.log_dict)
|
||||
self.log_dict = {}
|
||||
self.bbox_media_panel_images = []
|
||||
if self.result_artifact:
|
||||
self.result_artifact.add(self.result_table, 'result')
|
||||
wandb.log_artifact(self.result_artifact, aliases=['latest', 'last', 'epoch ' + str(self.current_epoch),
|
||||
('best' if best_result else '')])
|
||||
|
||||
wandb.log({"evaluation": self.result_table})
|
||||
self.result_table = wandb.Table(["epoch", "id", "ground truth", "prediction", "avg_confidence"])
|
||||
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
||||
|
||||
def finish_run(self):
|
||||
if self.wandb_run:
|
||||
if self.log_dict:
|
||||
with all_logging_disabled():
|
||||
wandb.log(self.log_dict)
|
||||
wandb.run.finish()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def all_logging_disabled(highest_level=logging.CRITICAL):
|
||||
""" source - https://gist.github.com/simon-weber/7853144
|
||||
A context manager that will prevent any logging messages triggered during the body from being processed.
|
||||
:param highest_level: the maximum logging level in use.
|
||||
This would only need to be changed if a custom level greater than CRITICAL is defined.
|
||||
"""
|
||||
previous_level = logging.root.manager.disable
|
||||
logging.disable(highest_level)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
logging.disable(previous_level)
|
||||
Reference in new issue
Block a user