Add train.py and val.py callbacks (#4220)
* added callbacks * Update callbacks.py * Update train.py * Update val.py * Fix CamlCase add staticmethod * Refactor logger into callbacks * Cleanup * New callback on_val_image_end() * Add curves and results images to TensorBoard Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Kalen MichaelandGlenn Jocher authored and GitHub committed
1 parent
d8f18834a2
commit
b74929c910
6 files changed
+230
-41
No files matched your search
@@ -0,0 +1,176 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
class Callbacks:
|
||||
""""
|
||||
Handles all registered callbacks for YOLOv5 Hooks
|
||||
"""
|
||||
|
||||
_callbacks = {
|
||||
'on_pretrain_routine_start': [],
|
||||
'on_pretrain_routine_end': [],
|
||||
|
||||
'on_train_start': [],
|
||||
'on_train_epoch_start': [],
|
||||
'on_train_batch_start': [],
|
||||
'optimizer_step': [],
|
||||
'on_before_zero_grad': [],
|
||||
'on_train_batch_end': [],
|
||||
'on_train_epoch_end': [],
|
||||
|
||||
'on_val_start': [],
|
||||
'on_val_batch_start': [],
|
||||
'on_val_image_end': [],
|
||||
'on_val_batch_end': [],
|
||||
'on_val_end': [],
|
||||
|
||||
'on_fit_epoch_end': [], # fit = train + val
|
||||
'on_model_save': [],
|
||||
'on_train_end': [],
|
||||
|
||||
'teardown': [],
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def register_action(self, hook, name='', callback=None):
|
||||
"""
|
||||
Register a new action to a callback hook
|
||||
|
||||
Args:
|
||||
hook The callback hook name to register the action to
|
||||
name The name of the action
|
||||
callback The callback to fire
|
||||
"""
|
||||
assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
|
||||
assert callable(callback), f"callback '{callback}' is not callable"
|
||||
self._callbacks[hook].append({'name': name, 'callback': callback})
|
||||
|
||||
def get_registered_actions(self, hook=None):
|
||||
""""
|
||||
Returns all the registered actions by callback hook
|
||||
|
||||
Args:
|
||||
hook The name of the hook to check, defaults to all
|
||||
"""
|
||||
if hook:
|
||||
return self._callbacks[hook]
|
||||
else:
|
||||
return self._callbacks
|
||||
|
||||
@staticmethod
|
||||
def run_callbacks(register, *args, **kwargs):
|
||||
"""
|
||||
Loop through the registered actions and fire all callbacks
|
||||
"""
|
||||
for logger in register:
|
||||
# print(f"Running callbacks.{logger['callback'].__name__}()")
|
||||
logger['callback'](*args, **kwargs)
|
||||
|
||||
def on_pretrain_routine_start(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the start of each pretraining routine
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_pretrain_routine_start'], *args, **kwargs)
|
||||
|
||||
def on_pretrain_routine_end(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the end of each pretraining routine
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_pretrain_routine_end'], *args, **kwargs)
|
||||
|
||||
def on_train_start(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the start of each training
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_train_start'], *args, **kwargs)
|
||||
|
||||
def on_train_epoch_start(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the start of each training epoch
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_train_epoch_start'], *args, **kwargs)
|
||||
|
||||
def on_train_batch_start(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the start of each training batch
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_train_batch_start'], *args, **kwargs)
|
||||
|
||||
def optimizer_step(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks on each optimizer step
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['optimizer_step'], *args, **kwargs)
|
||||
|
||||
def on_before_zero_grad(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks before zero grad
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_before_zero_grad'], *args, **kwargs)
|
||||
|
||||
def on_train_batch_end(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the end of each training batch
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_train_batch_end'], *args, **kwargs)
|
||||
|
||||
def on_train_epoch_end(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the end of each training epoch
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_train_epoch_end'], *args, **kwargs)
|
||||
|
||||
def on_val_start(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the start of the validation
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_val_start'], *args, **kwargs)
|
||||
|
||||
def on_val_batch_start(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the start of each validation batch
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_val_batch_start'], *args, **kwargs)
|
||||
|
||||
def on_val_image_end(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the end of each val image
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_val_image_end'], *args, **kwargs)
|
||||
|
||||
def on_val_batch_end(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the end of each validation batch
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_val_batch_end'], *args, **kwargs)
|
||||
|
||||
def on_val_end(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the end of the validation
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_val_end'], *args, **kwargs)
|
||||
|
||||
def on_fit_epoch_end(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the end of each fit (train+val) epoch
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_fit_epoch_end'], *args, **kwargs)
|
||||
|
||||
def on_model_save(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks after each model save
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_model_save'], *args, **kwargs)
|
||||
|
||||
def on_train_end(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks at the end of training
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['on_train_end'], *args, **kwargs)
|
||||
|
||||
def teardown(self, *args, **kwargs):
|
||||
"""
|
||||
Fires all registered callbacks before teardown
|
||||
"""
|
||||
self.run_callbacks(self._callbacks['teardown'], *args, **kwargs)
|
||||
@@ -67,6 +67,11 @@ def try_except(func):
|
||||
return handler
|
||||
|
||||
|
||||
def methods(instance):
|
||||
# Get class/instance methods
|
||||
return [f for f in dir(instance) if callable(getattr(instance, f)) and not f.startswith("__")]
|
||||
|
||||
|
||||
def set_logging(rank=-1, verbose=True):
|
||||
logging.basicConfig(
|
||||
format="%(message)s",
|
||||
|
||||
+24
-21
@@ -29,10 +29,12 @@ class Loggers():
|
||||
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
|
||||
|
||||
def start(self):
|
||||
self.csv = True # always log to csv
|
||||
|
||||
# Message
|
||||
@@ -57,7 +59,11 @@ class Loggers():
|
||||
else:
|
||||
self.wandb = None
|
||||
|
||||
return self
|
||||
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):
|
||||
# Callback runs on train batch end
|
||||
@@ -78,8 +84,8 @@ class Loggers():
|
||||
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
|
||||
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)
|
||||
|
||||
@@ -89,25 +95,20 @@ class Loggers():
|
||||
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
|
||||
def on_fit_epoch_end(self, mloss, results, lr, epoch, best_fitness, fi):
|
||||
# Callback runs at the end of each fit (train+val) epoch
|
||||
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
|
||||
|
||||
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'] + keys)).rstrip(',') + '\n') # add header
|
||||
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) # TensorBoard
|
||||
self.tb.add_scalar(k, v, epoch)
|
||||
|
||||
if self.wandb:
|
||||
self.wandb.log(x)
|
||||
@@ -119,20 +120,22 @@ class Loggers():
|
||||
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):
|
||||
def on_train_end(self, last, best, plots, epoch):
|
||||
# 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.tb:
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
for f in files:
|
||||
self.tb.add_image(f.stem, np.asarray(Image.open(f)), epoch, dataformats='HWC')
|
||||
|
||||
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]})
|
||||
+1
-5
@@ -281,7 +281,7 @@ def plot_study_txt(path='', x=None): # from utils.plots import *; plot_study_tx
|
||||
plt.savefig(str(Path(path).name) + '.png', dpi=300)
|
||||
|
||||
|
||||
def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
||||
def plot_labels(labels, names=(), save_dir=Path('')):
|
||||
# plot dataset labels
|
||||
print('Plotting labels... ')
|
||||
c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes
|
||||
@@ -324,10 +324,6 @@ def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
||||
matplotlib.use('Agg')
|
||||
plt.close()
|
||||
|
||||
# loggers
|
||||
if loggers:
|
||||
loggers.log_images(save_dir.glob('*labels*.jpg'))
|
||||
|
||||
|
||||
def plot_evolution(yaml_file='data/hyp.finetune.yaml'): # from utils.plots import *; plot_evolution()
|
||||
# Plot hyperparameter evolution results in evolve.txt
|
||||
|
||||
Reference in new issue
Block a user