Spaces:
Running
Running
| # YOLOv5 🚀 by Ultralytics, GPL-3.0 license | |
| """ | |
| Logging utils | |
| """ | |
| import os | |
| import warnings | |
| from threading import Thread | |
| import pkg_resources as pkg | |
| import torch | |
| from torch.utils.tensorboard import SummaryWriter | |
| from utils.general import colorstr, cv2, emojis | |
| from utils.loggers.wandb.wandb_utils import WandbLogger | |
| from utils.plots import plot_images, plot_results | |
| from utils.torch_utils import de_parallel | |
| LOGGERS = ('csv', 'tb', 'wandb') # text-file, TensorBoard, Weights & Biases | |
| RANK = int(os.getenv('RANK', -1)) | |
| try: | |
| import wandb | |
| assert hasattr(wandb, '__version__') # verify package import not local dir | |
| if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in [0, -1]: | |
| try: | |
| wandb_login_success = wandb.login(timeout=30) | |
| except wandb.errors.UsageError: # known non-TTY terminal issue | |
| wandb_login_success = False | |
| if not wandb_login_success: | |
| wandb = None | |
| except (ImportError, AssertionError): | |
| wandb = None | |
| 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 | |
| self.best_keys = ['best/epoch', 'best/precision', 'best/recall', 'best/mAP_0.5', 'best/mAP_0.5:0.95'] | |
| 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)" | |
| self.logger.info(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) | |
| # temp warn. because nested artifacts not supported after 0.12.10 | |
| if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.11'): | |
| self.logger.warning( | |
| "YOLOv5 temporarily requires wandb version 0.12.10 or below. Some features may not work as expected." | |
| ) | |
| else: | |
| self.wandb = None | |
| def on_train_start(self): | |
| # Callback runs on train start | |
| pass | |
| 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 | |
| if plots: | |
| if ni == 0: | |
| if not self.opt.sync_bn: # --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: | |
| if best_fitness == fi: | |
| best_results = [epoch] + vals[3:7] | |
| for i, name in enumerate(self.best_keys): | |
| self.wandb.wandb_run.summary[name] = best_results[i] # log best results in the summary | |
| 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: | |
| for f in files: | |
| self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC') | |
| if self.wandb: | |
| self.wandb.log({k: v for k, v in zip(self.keys[3:10], results)}) # log best.pt val results | |
| 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() | |
| 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) | |