fancier and more flexible (self adapting to loss_dict) console logging. Fixing multi-gpu loss reduce

utils/generic_utils.py

@@ -144,7 +144,7 @@ def save_best_model(model, optimizer, model_loss, best_loss, out_path,
         best_loss = model_loss
         bestmodel_path = 'best_model.pth.tar'
         bestmodel_path = os.path.join(out_path, bestmodel_path)
-        print("\n > BEST MODEL ({0:.5f}) : {1:}".format(
+        print(" > BEST MODEL ({0:.5f}) : {1:}".format(
             model_loss, bestmodel_path))
         torch.save(state, bestmodel_path)
     return best_loss
@@ -368,6 +368,9 @@ class KeepAverage():
     def __getitem__(self, key):
         return self.avg_values[key]
 
+    def items(self):
+        return self.avg_values.items() 
+
     def add_value(self, name, init_val=0, init_iter=0):
         self.avg_values[name] = init_val
         self.iters[name] = init_iter
@@ -407,6 +410,37 @@ def _check_argument(name, c, enum_list=None, max_val=None, min_val=None, restric
             assert isinstance(c[name], val_type) or c[name] is None, f' [!] {name} has wrong type - {type(c[name])} vs {val_type}'
 
 
+tcolors = AttrDict({
+    'OKBLUE': '\033[94m',
+    'HEADER': '\033[95m', 
+    'OKGREEN': '\033[92m',
+    'WARNING': '\033[93m',
+    'FAIL': '\033[91m',
+    'ENDC': '\033[0m',
+    'BOLD': '\033[1m',
+    'UNDERLINE': '\033[4m'
+})
+
+
+def print_train_step(batch_steps, step, global_step, avg_spec_length, avg_text_length, step_time, loader_time, lr, print_dict):
+    indent = "     | > "
+    print()
+    log_text = "{}   --> STEP: {}/{} -- GLOBAL_STEP: {}{}\n".format(tcolors.BOLD, step, batch_steps, global_step, tcolors.ENDC)
+    for key, value in print_dict.items():
+        log_text += "{}{}: {:.5f}\n".format(indent, key, value)
+    log_text += "{}avg_spec_len: {}\n{}avg_text_len: {}\n{}step_time: {:.2f}\n{}loader_time: {:.2f}\n{}lr: {:.5f}"\
+        .format(indent, avg_spec_length, indent, avg_text_length, indent, step_time, indent, loader_time, indent, lr)
+    print(log_text, flush=True)
+    
+
+def print_train_epoch(step, global_step, epoch, loss_dict):
+    pass
+
+
+def print_eval_step():
+    pass
+
+
 def check_config(c):
     _check_argument('model', c, enum_list=['tacotron', 'tacotron2'], restricted=True, val_type=str)
     _check_argument('run_name', c, restricted=True, val_type=str)