Fix total time

laiguokun · Jan 21, 2017 · c09e38b · c09e38b
1 parent 8c948e2
commit c09e38b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/main.py b/main.py
@@ -1,4 +1,3 @@
-
 '''Train CIFAR10 with PyTorch.'''
 from __future__ import print_function
 
@@ -31,7 +30,6 @@
 
 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
 
-
 # Model.
 print('==> Building model..')
 class Net(nn.Module):

diff --git a/utils.py b/utils.py
@@ -40,11 +40,14 @@ def get_mean_and_std(dataset):
 last_time = time.time()
 begin_time = last_time
 def progress_bar(current, total, msg=None):
-    global last_time
-    sys.stdout.write(' [')
+    global last_time, begin_time
+    if current == 0:
+        begin_time = time.time()  # Reset for new bar.
+
     cur_len = int(TOTAL_BAR_LENGTH*current/total)
     rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1
 
+    sys.stdout.write(' [')
     for i in range(cur_len):
         sys.stdout.write('=')
     sys.stdout.write('>')
@@ -56,6 +59,7 @@ def progress_bar(current, total, msg=None):
     cur_time = time.time()
     step_time = cur_time - last_time
     last_time = cur_time
+
     tot_time = cur_time - begin_time
     L.append('  Step: %s' % format_time(step_time))
     L.append(' | Tot: %s' % format_time(tot_time))