added progbar

tensorgraph/progbar.py

@@ -0,0 +1,87 @@
+from __future__ import absolute_import
+import numpy as np
+import time
+import sys
+
+
+class ProgressBar(object):
+    def __init__(self, target, width=30, verbose=1):
+        '''
+            @param target: total number of steps expected
+        '''
+        self.width = width
+        self.target = target
+        self.sum_values = {}
+        self.unique_values = []
+        self.start = time.time()
+        self.total_width = 0
+        self.seen_so_far = 0
+        self.verbose = verbose
+
+    def update(self, current, values=[]):
+        '''
+            @param current: index of current step
+            @param values: list of tuples (name, value_for_last_step).
+            The progress bar will display averages for these values.
+        '''
+        for k, v in values:
+            if k not in self.sum_values:
+                self.sum_values[k] = [v * (current-self.seen_so_far), current-self.seen_so_far]
+                self.unique_values.append(k)
+            else:
+                self.sum_values[k][0] += v * (current-self.seen_so_far)
+                self.sum_values[k][1] += (current-self.seen_so_far)
+        self.seen_so_far = current
+
+        now = time.time()
+        if self.verbose == 1:
+            prev_total_width = self.total_width
+            sys.stdout.write("\b" * prev_total_width)
+            sys.stdout.write("\r")
+
+            numdigits = int(np.floor(np.log10(self.target))) + 1
+            barstr = '%%%dd/%%%dd [' % (numdigits, numdigits)
+            bar = barstr % (current, self.target)
+            prog = float(current)/self.target
+            if current > self.target:
+                prog = 1
+            prog_width = int(self.width*prog)
+            if prog_width > 0:
+                bar += ('='*(prog_width-1))
+                if current < self.target:
+                    bar += '>'
+            bar += ('.'*(self.width-prog_width-1))
+            bar += ']'
+            sys.stdout.write(bar)
+            self.total_width = len(bar)
+
+            if current:
+                time_per_unit = (now - self.start) / current
+            else:
+                time_per_unit = 0
+            eta = time_per_unit*(self.target - current)
+            info = ''
+            if current < self.target:
+                info += ' - ETA: %ds' % eta
+            else:
+                info += ' - %ds' % (now - self.start)
+            for k in self.unique_values:
+                info += ' - %s: %.4f' % (k, self.sum_values[k][0]/ max(1, self.sum_values[k][1]))
+
+            self.total_width += len(info)
+            if prev_total_width > self.total_width:
+                info += ((prev_total_width-self.total_width) * " ")
+
+            sys.stdout.write(info)
+            sys.stdout.flush()
+
+        if self.verbose == 2:
+            if current >= self.target:
+                info = '%ds' % (now - self.start)
+                for k in self.unique_values:
+                    info += ' - %s: %.4f' % (k, self.sum_values[k][0]/ max(1, self.sum_values[k][1]))
+                sys.stdout.write(info + "\n")
+
+
+    def add(self, n, values=[]):
+        self.update(self.seen_so_far+n, values)