New CSV Logger (#4148)

* New CSV Logger

* cleanup

* move batch plots into Logger

* rename comment

* Remove total loss from progress bar

* mloss :-1 bug fix

* Update plot_results()

* Update plot_results()

* plot_results bug fix

utils/loggers/__init__.py

@@ -1,15 +1,17 @@
 # YOLOv5 experiment logging utils
 
 import warnings
+from threading import Thread
 
 import torch
 from torch.utils.tensorboard import SummaryWriter
 
 from utils.general import colorstr, 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 = ('txt', 'tb', 'wandb')  # text-file, TensorBoard, Weights & Biases
+LOGGERS = ('csv', 'tb', 'wandb')  # text-file, TensorBoard, Weights & Biases
 
 try:
     import wandb
@@ -21,10 +23,8 @@ except (ImportError, AssertionError):
 
 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):
+    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.results_file = results_file
         self.weights = weights
         self.opt = opt
         self.hyp = hyp
@@ -35,7 +35,7 @@ class Loggers():
             setattr(self, k, None)  # init empty logger dictionary
 
     def start(self):
-        self.txt = True  # always log to txt
+        self.csv = True  # always log to csv
 
         # Message
         try:
@@ -63,15 +63,19 @@ class Loggers():
 
         return self
 
-    def on_train_batch_end(self, ni, model, imgs):
+    def on_train_batch_end(self, ni, model, imgs, targets, paths, plots):
         # 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()]})
+        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
@@ -89,21 +93,28 @@ 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, 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',
+    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
-        if self.txt:
-            with open(self.results_file, 'a') as f:
-                f.write(s + '%10.4g' * 7 % results + '\n')  # append metrics, val_loss
+        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 x, tag in zip(vals, tags):
-                self.tb.add_scalar(tag, x, epoch)  # TensorBoard
+            for k, v in x.items():
+                self.tb.add_scalar(k, v, epoch)  # TensorBoard
+
         if self.wandb:
-            self.wandb.log({k: v for k, v in zip(tags, vals)})
+            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):
@@ -112,8 +123,10 @@ 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):
+    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: