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

utils/loggers/__init__.py

@@ -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]})

utils/loggers/wandb/sweep.py

@@ -1,12 +1,12 @@
 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
-import test
 from utils.general import increment_path
 from utils.torch_utils import select_device
 

utils/loggers/wandb/sweep.yaml

@@ -14,7 +14,7 @@
 # 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/wandb_logging/sweep.py
+program: utils/loggers/wandb/sweep.py
 method: random
 metric:
   name: metrics/mAP_0.5

utils/loggers/wandb/wandb_utils.py

@@ -1,4 +1,5 @@
 """Utilities and tools for tracking runs with Weights & Biases."""
+
 import logging
 import os
 import sys
@@ -8,15 +9,18 @@ from pathlib import Path
 import yaml
 from tqdm import tqdm
 
-sys.path.append(str(Path(__file__).parent.parent.parent))  # add utils/ to path
+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 colorstr, check_dataset, check_file
+from utils.general import check_dataset, check_file
 
 try:
     import wandb
-    from wandb import init, finish
-except ImportError:
+
+    assert hasattr(wandb, '__version__')  # verify package import not local dir
+except (ImportError, AssertionError):
     wandb = None
 
 RANK = int(os.getenv('RANK', -1))
@@ -106,7 +110,7 @@ class WandbLogger():
         self.data_dict = data_dict
         self.bbox_media_panel_images = []
         self.val_table_path_map = None
-        self.max_imgs_to_log = 16 
+        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):
@@ -134,13 +138,11 @@ class WandbLogger():
                 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.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)
-        else:
-            prefix = colorstr('wandb: ')
-            print(f"{prefix}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
 
     def check_and_upload_dataset(self, opt):
         assert wandb, 'Install wandb to upload dataset'
@@ -169,7 +171,7 @@ class WandbLogger():
                                                                                            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)
@@ -177,7 +179,6 @@ class WandbLogger():
             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"])
@@ -315,9 +316,9 @@ class WandbLogger():
                                    )
 
     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
+        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
+        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]},
@@ -328,7 +329,6 @@ class WandbLogger():
                     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():

utils/plots.py

@@ -327,9 +327,8 @@ def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
     plt.close()
 
     # loggers
-    for k, v in loggers.items() or {}:
-        if k == 'wandb' and v:
-            v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]}, commit=False)
+    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()