right allign the log

open-mmlab · Aug 15, 2022 · 9a01a9c · 9a01a9c
1 parent 7b3397c
commit 9a01a9c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/mmengine/logging/log_processor.py b/mmengine/logging/log_processor.py
@@ -141,20 +141,40 @@ def get_log_after_iter(self, runner, batch_idx: int,
         #  train: Epoch [5/10000] ... (divided by `max_iter`)
         #  val/test: Epoch [5/2000] ... (divided by length of dataloader)
         if self.by_epoch:
+            # Align the iteration log:
+            # Epoch(train)  [9][010/270]
+            # ...               ||| |||
+            # Epoch(train)  [9][100/270]
+            max_dataloader_len = len(current_loop.dataloader)
+            max_str_len = len(str(max_dataloader_len))
+            cur_iter_str = str(cur_iter).rjust(max_str_len, '0')
+
             if mode in ['train', 'val']:
+                # Right Align the epoch log:
+                # Epoch(train)   [9][100/270]
+                # ...             ||
+                # Epoch(train) [100][100/270]
                 cur_epoch = self._get_epoch(runner, mode)
+                max_epochs = runner.max_epochs
+                max_str_len = len(str(max_epochs)) + 3
+                cur_epoch_str = str(f'[{cur_epoch}]').rjust(max_str_len, ' ')
                 tag['epoch'] = cur_epoch
-                log_str = (f'Epoch({mode}) [{cur_epoch}]'
-                           f'[{cur_iter}/{len(current_loop.dataloader)}]  ')
+                log_str = (f'Epoch({mode}){cur_epoch_str}'
+                           f'[{cur_iter_str}/{max_dataloader_len}]  ')
             else:
                 log_str = (f'Epoch({mode}) '
-                           f'[{cur_iter}/{len(current_loop.dataloader)}]  ')
+                           f'[{cur_iter_str}/{max_dataloader_len}]  ')
         else:
             if mode == 'train':
+                max_str_len = len(str(runner.max_iters))
+                cur_iter_str = str(cur_iter).rjust(max_str_len, '0')
                 log_str = (f'Iter({mode}) '
-                           f'[{cur_iter}/{runner.max_iters}]  ')
+                           f'[{cur_iter_str}/{runner.max_iters}]  ')
             else:
-                log_str = (f'Iter({mode}) [{batch_idx+1}'
+                max_dataloader_len = len(current_loop.dataloader)
+                max_str_len = len(str(max_dataloader_len))
+                cur_iter_str = str(batch_idx + 1).rjust(max_str_len, '0')
+                log_str = (f'Iter({mode}) [{cur_iter_str}'
                            f'/{len(current_loop.dataloader)}]  ')
         # Concatenate lr, momentum string with log header.
         log_str += f'{lr_str}  '

diff --git a/tests/test_logging/test_log_processor.py b/tests/test_logging/test_log_processor.py
@@ -90,7 +90,7 @@ def test_get_log_after_iter(self, by_epoch, mode):
         if by_epoch:
             if mode in ['train', 'val']:
                 cur_epoch = log_processor._get_epoch(self.runner, mode)
-                log_str = (f'Epoch({mode}) [{cur_epoch}][2/'
+                log_str = (f'Epoch({mode})  [{cur_epoch}][02/'
                            f'{len(cur_loop.dataloader)}]  ')
             else:
                 log_str = (f'Epoch({mode}) [2/{len(cur_loop.dataloader)}]  ')
@@ -113,6 +113,9 @@ def test_get_log_after_iter(self, by_epoch, mode):
             if mode == 'train':
                 max_iters = self.runner.max_iters
                 log_str = f'Iter({mode}) [11/{max_iters}]  '
+            elif mode == 'val':
+                max_iters = len(cur_loop.dataloader)
+                log_str = f'Iter({mode}) [02/{max_iters}]  '
             else:
                 max_iters = len(cur_loop.dataloader)
                 log_str = f'Iter({mode}) [2/{max_iters}]  '
@@ -237,6 +240,7 @@ def test_get_cur_loop(self):
     def setup(self):
         runner = MagicMock()
         runner.epoch = 1
+        runner.max_epochs = 10
         runner.iter = 10
         runner.max_iters = 50
         runner.train_dataloader = [0] * 20