Fixed progress bar display issue and no error printing issue

main.py

@@ -5,6 +5,7 @@
 import config
 import os
 import sys
+import traceback
 import time
 import datetime
 import tensorflow as tf
@@ -17,7 +18,7 @@
 from checkpoints import Checkpoints
 
 def main():
-    # parse the arguments
+    # Parse the Arguments
     args = config.parse_args()
     random.seed(args.manual_seed)
     tf.set_random_seed(args.manual_seed)
@@ -27,36 +28,32 @@ def main():
     if args.save_results:
         utils.saveargs(args)
 
-    # initialize the checkpoint class
+    # Initialize the Checkpoints Class
     checkpoints = Checkpoints(args)
 
     # Create Model
     models = Model(args)
     model, criterion, evaluation = models.setup(checkpoints)
 
-    # initialize a sample input to build the model for the first time and print its summary
-    batch = tf.zeros((1, args.nchannels, args.resolution_high, args.resolution_wide))
-    model(batch)
+    # Print Model Summary
     print('Model summary: {}'.format(model.name))
     print(model.summary())
 
     # Data Loading
     dataloader_obj = Dataloader(args)
     dataloader = dataloader_obj.create()
 
-    # initialize trainer and tester
+    # Initialize Trainer and Tester
     trainer = Trainer(args, model, criterion, evaluation)
     tester = Tester(args, model, criterion, evaluation)
 
-    # start training !!!
+    # Start Training !!!
     loss_best = 1e10
     for epoch in range(args.nepochs):
-        print('\nEpoch %d/%d\n' % (epoch + 1, args.nepochs))
+        print('\nEpoch %d/%d' % (epoch + 1, args.nepochs))
 
-        # train for a single epoch
-        print("Training...")
+        # Train and Test for a Single Epoch
         loss_train = trainer.train(epoch, dataloader["train"])
-        print("Testing...")
         loss_test = tester.test(epoch, dataloader["test"])
 
         if loss_best > loss_test:

model.py

@@ -4,6 +4,7 @@
 import models
 import losses
 import evaluate
+import tensorflow as tf
 
 class Model:
     def __init__(self, args):
@@ -26,4 +27,8 @@ def setup(self, checkpoints):
         else:
             model = checkpoints.load(model, ckpt)
 
+        # initialize a sample input to build the model for the first time
+        batch = tf.zeros((1, self.args.nchannels, self.args.resolution_high, self.args.resolution_wide))
+        model(batch)
+
         return model, criterion, evaluation

test.py

@@ -66,7 +66,7 @@ def test(self, epoch, dataloader):
         if self.log_type == 'progressbar':
             # Progress bar
             processed_data_len = 0
-            bar = plugins.Bar('{:<10}'.format('Train'), max=data_len)
+            bar = plugins.Bar('{:<10}'.format('Test'), max=data_len//self.batch_size)
         end = time.time()
 
         with self.summary_writer.as_default():

train.py

@@ -70,7 +70,7 @@ def train(self, epoch, dataloader):
         if self.log_type == 'progressbar':
             # Progress bar
             processed_data_len = 0
-            bar = plugins.Bar('{:<10}'.format('Train'), max=data_len)
+            bar = plugins.Bar('{:<10}'.format('Train'), max=data_len//self.batch_size)
         end = time.time()
 
         with self.summary_writer.as_default():

utils.py

@@ -27,7 +27,7 @@ def cleanup():
         except OSError as ex:
             raise Exception("wasn't able to kill the child process (pid:{}).".format(child.pid))
     #     # os.waitpid(child.pid, os.P_ALL)
-    print('\x1b[?25h', end='', flush=True)  # show cursor
+    print('\n\n\x1b[?25h', end='', flush=True)  # show cursor
     sys.exit(0)