[WIP] fix progress bar lacks a wrap

mmcv/engine/test.py

@@ -27,17 +27,17 @@ def single_gpu_test(model, data_loader):
     model.eval()
     results = []
     dataset = data_loader.dataset
-    prog_bar = mmcv.ProgressBar(len(dataset))
-    for data in data_loader:
-        with torch.no_grad():
-            result = model(return_loss=False, **data)
-        results.extend(result)
-
-        # Assume result has the same length of batch_size
-        # refer to https://github.com/open-mmlab/mmcv/issues/985
-        batch_size = len(result)
-        for _ in range(batch_size):
-            prog_bar.update()
+    with mmcv.ProgressBar(len(dataset)) as pb:
+        for data in data_loader:
+            with torch.no_grad():
+                result = model(return_loss=False, **data)
+            results.extend(result)
+
+            # Assume result has the same length of batch_size
+            # refer to https://github.com/open-mmlab/mmcv/issues/985
+            batch_size = len(result)
+            for _ in range(batch_size):
+                pb.update()
     return results
 
 
@@ -64,21 +64,23 @@ def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False):
     results = []
     dataset = data_loader.dataset
     rank, world_size = get_dist_info()
-    if rank == 0:
-        prog_bar = mmcv.ProgressBar(len(dataset))
     time.sleep(2)  # This line can prevent deadlock problem in some cases.
-    for i, data in enumerate(data_loader):
-        with torch.no_grad():
-            result = model(return_loss=False, **data)
-        results.extend(result)
-
-        if rank == 0:
-            batch_size = len(result)
-            batch_size_all = batch_size * world_size
-            if batch_size_all + prog_bar.completed > len(dataset):
-                batch_size_all = len(dataset) - prog_bar.completed
-            for _ in range(batch_size_all):
-                prog_bar.update()
+    # The test and progress bar are coupled, so we instantiate progress bar
+    # for all process. But we hope progress bar only displays on the master
+    # process, so we use `start=(rank == 0)`.
+    with mmcv.ProgressBar(len(dataset), start=(rank == 0)) as pb:
+        for i, data in enumerate(data_loader):
+            with torch.no_grad():
+                result = model(return_loss=False, **data)
+            results.extend(result)
+
+            if rank == 0:
+                batch_size = len(result)
+                batch_size_all = batch_size * world_size
+                if batch_size_all + pb.completed > len(dataset):
+                    batch_size_all = len(dataset) - pb.completed
+                for _ in range(batch_size_all):
+                    pb.update()
 
     # collect results from all ranks
     if gpu_collect:

mmcv/runner/hooks/evaluation.py

@@ -500,7 +500,6 @@ def _do_evaluate(self, runner):
             tmpdir=tmpdir,
             gpu_collect=self.gpu_collect)
         if runner.rank == 0:
-            print('\n')
             runner.log_buffer.output['eval_iter_num'] = len(self.dataloader)
             key_score = self.evaluate(runner, results)
             # the key_score may be `None` so it needs to skip the action to

mmcv/utils/progressbar.py

@@ -18,6 +18,12 @@ def __init__(self, task_num=0, bar_width=50, start=True, file=sys.stdout):
         if start:
             self.start()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.file.write('\n')
+
     @property
     def terminal_width(self):
         width, _ = get_terminal_size()
@@ -86,12 +92,11 @@ def track_progress(func, tasks, bar_width=50, file=sys.stdout, **kwargs):
     else:
         raise TypeError(
             '"tasks" must be an iterable object or a (iterator, int) tuple')
-    prog_bar = ProgressBar(task_num, bar_width, file=file)
-    results = []
-    for task in tasks:
-        results.append(func(task, **kwargs))
-        prog_bar.update()
-    prog_bar.file.write('\n')
+    with ProgressBar(task_num, bar_width, file=file) as pb:
+        results = []
+        for task in tasks:
+            results.append(func(task, **kwargs))
+            pb.update()
     return results
 
 
@@ -155,22 +160,21 @@ def track_parallel_progress(func,
     pool = init_pool(nproc, initializer, initargs)
     start = not skip_first
     task_num -= nproc * chunksize * int(skip_first)
-    prog_bar = ProgressBar(task_num, bar_width, start, file=file)
-    results = []
-    if keep_order:
-        gen = pool.imap(func, tasks, chunksize)
-    else:
-        gen = pool.imap_unordered(func, tasks, chunksize)
-    for result in gen:
-        results.append(result)
-        if skip_first:
-            if len(results) < nproc * chunksize:
-                continue
-            elif len(results) == nproc * chunksize:
-                prog_bar.start()
-                continue
-        prog_bar.update()
-    prog_bar.file.write('\n')
+    with ProgressBar(task_num, bar_width, start, file=file) as pb:
+        results = []
+        if keep_order:
+            gen = pool.imap(func, tasks, chunksize)
+        else:
+            gen = pool.imap_unordered(func, tasks, chunksize)
+        for result in gen:
+            results.append(result)
+            if skip_first:
+                if len(results) < nproc * chunksize:
+                    continue
+                elif len(results) == nproc * chunksize:
+                    pb.start()
+                    continue
+            pb.update()
     pool.close()
     pool.join()
     return results
@@ -201,8 +205,7 @@ def track_iter_progress(tasks, bar_width=50, file=sys.stdout):
     else:
         raise TypeError(
             '"tasks" must be an iterable object or a (iterator, int) tuple')
-    prog_bar = ProgressBar(task_num, bar_width, file=file)
-    for task in tasks:
-        yield task
-        prog_bar.update()
-    prog_bar.file.write('\n')
+    with ProgressBar(task_num, bar_width, file=file) as pb:
+        for task in tasks:
+            yield task
+            pb.update()