class Loggers():
# YOLOv5 Loggers class
def __init__(self,
save_dir=None,
weights=None,
opt=None,
hyp=None,
logger=None,
include=LOGGERS):
self.save_dir = save_dir
self.weights = weights
self.opt = opt
self.hyp = hyp
self.logger = logger # for printing results to console
self.include = include
self.keys = [
'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', # metrics
'val/box_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
for k in LOGGERS:
setattr(self, k, None) # init empty logger dictionary
self.csv = True # always log to csv
# Message
if not wandb:
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
if wandb and 'wandb' in self.include:
wandb_artifact_resume = isinstance(
self.opt.resume,
str) and self.opt.resume.startswith('wandb-artifact://')
run_id = torch.load(self.weights).get(
'wandb_id'
) if self.opt.resume and not wandb_artifact_resume else None
self.opt.hyp = self.hyp # add hyperparameters
self.wandb = WandbLogger(self.opt, run_id)
else:
self.wandb = None
def on_pretrain_routine_end(self):
# Callback runs on pre-train routine end
paths = self.save_dir.glob('*labels*.jpg') # training labels
if self.wandb:
self.wandb.log({
"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]
})
def on_train_batch_end(self, ni, model, imgs, targets, paths, plots,
sync_bn):
# Callback runs on train batch end
if plots:
if ni == 0:
if not sync_bn: # tb.add_graph() --sync known issue https://github.com/ultralytics/yolov5/issues/3754
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 ni < 3:
f = self.save_dir / f'train_batch{ni}.jpg' # filename
Thread(target=plot_images,
args=(imgs, targets, paths, f),
daemon=True).start()
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_image_end(self, pred, predn, path, names, im):
# Callback runs on val image 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_fit_epoch_end(self, vals, epoch, best_fitness, fi):
# Callback runs at the end of each fit (train+val) epoch
x = {k: v for k, v in zip(self.keys, vals)} # dict
if self.csv:
file = self.save_dir / 'results.csv'
n = len(x) + 1 # number of cols
s = '' if file.exists() else (
('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') +
'\n') # add header
with open(file, 'a') as f:
f.write(s +
('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') +
'\n')
if self.tb:
for k, v in x.items():
self.tb.add_scalar(k, v, epoch)
if self.wandb:
self.wandb.log(x)
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, plots, epoch, results):
# Callback runs on training end
if plots:
plot_results(file=self.save_dir /
'results.csv') # save results.png
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.tb:
import cv2
for f in files:
self.tb.add_image(f.stem,
cv2.imread(str(f))[..., ::-1],
epoch,
dataformats='HWC')
if self.wandb:
self.wandb.log({
"Results":
[wandb.Image(str(f), caption=f.name) for f in files]
})
# Calling wandb.log. TODO: Refactor this into WandbLogger.log_model
if not self.opt.evolve:
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()
else:
self.wandb.finish_run()
self.wandb = WandbLogger(self.opt)
def on_params_update(self, params):
# Update hyperparams or configs of the experiment
# params: A dict containing {param: value} pairs
if self.wandb:
self.wandb.wandb_run.config.update(params, allow_val_change=True)
class Loggers():
# YOLOv5 Loggers class
def __init__(self,
save_dir=None,
weights=None,
opt=None,
hyp=None,
data_dict=None,
logger=None,
include=LOGGERS):
self.save_dir = save_dir
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.csv = True # always log to csv
# 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, run_id, self.data_dict)
except:
self.wandb = None
return self
def on_train_batch_end(self, ni, model, imgs, targets, paths, plots):
# Callback runs on train batch end
if plots:
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 ni < 3:
f = self.save_dir / f'train_batch{ni}.jpg' # filename
Thread(target=plot_images,
args=(imgs, targets, paths, f),
daemon=True).start()
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, best_fitness, fi):
# Callback runs on val end during training
vals = list(mloss) + list(results) + lr
keys = [
'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', # metrics
'val/box_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
x = {k: v for k, v in zip(keys, vals)} # dict
if self.csv:
file = self.save_dir / 'results.csv'
n = len(x) + 1 # number of cols
s = '' if file.exists() else (
('%20s,' * n % tuple(['epoch'] + keys)).rstrip(',') +
'\n') # add header
with open(file, 'a') as f:
f.write(s +
('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') +
'\n')
if self.tb:
for k, v in x.items():
self.tb.add_scalar(k, v, epoch) # TensorBoard
if self.wandb:
self.wandb.log(x)
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, plots):
# Callback runs on training end
if plots:
plot_results(dir=self.save_dir) # save results.png
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]
})