Repository: espnet/espnet
Issue: espnet2/train/reporter.py causes exception if  accum_grad is not 1

If accum_grad is not 1,  the number of elements of some keys such as 'optim_step_time' are not equal to the count value.
So, it causes the exception at line 175 of reporter.py

Here is the sample. 
when accum_grad = 2, self.count=2, but the number of 'optim_step_time' items is 1.

defaultdict(<class 'list'>, {'iter_time': [Average(value=0.1262605399970198), Average(value=8.855300256982446e-05)], 'forward_time': [Average(value=0.36080345801019575), Average(value=0.0742439620080404)], 'loss': [WeightedAverage(value=98.05116271972656, weight=64), WeightedAverage(value=327.82421875, weight=64)], 'loss_att': [WeightedAverage(value=27.909347534179688, weight=64), WeightedAverage(value=108.47477722167969, weight=64)], 'loss_ctc': [WeightedAverage(value=261.71539306640625, weight=64), WeightedAverage(value=839.6395263671875, weight=64)], 'acc': [WeightedAverage(value=0.0, weight=64), WeightedAverage(value=0.0, weight=64)], 'backward_time': [Average(value=0.22204561799298972), Average(value=0.20624846100690775)], 'optim_step_time': [Average(value=0.03714967500127386)], 'lr_0': [Average(value=2.2222222222222224e-07)], 'train_time': [Average(value=1.294450648012571)]})


Changed files:
espnet2/train/trainer.py, test/espnet2/train/test_reporter.py, espnet2/train/reporter.py

diff --git a/espnet2/train/reporter.py b/espnet2/train/reporter.py
index 56119f26b..b6871e5a7 100644
--- a/espnet2/train/reporter.py
+++ b/espnet2/train/reporter.py
@@ -126,6 +126,7 @@ class SubReporter:
         self._finished = False
         self.total_count = total_count
         self.count = 0
+        self._seen_keys_in_the_step = set()
 
     def get_total_count(self) -> int:
         """Returns the number of iterations over all epochs."""
@@ -134,26 +135,51 @@ class SubReporter:
     def get_epoch(self) -> int:
         return self.epoch
 
+    def next(self):
+        for key, stats_list in self.stats.items():
+            if key not in self._seen_keys_in_the_step:
+                # Fill nan value if the key is not registered in this step
+                if isinstance(stats_list[0], WeightedAverage):
+                    stats_list.append(to_reported_value(np.nan, 0))
+                elif isinstance(stats_list[0], Average):
+                    stats_list.append(to_reported_value(np.nan))
+                else:
+                    raise NotImplementedError(f"type={type(stats_list[0])}")
+
+            assert len(stats_list) == self.count, (len(stats_list), self.count)
+
+        self._seen_keys_in_the_step = set()
+
     def register(
         self,
         stats: Dict[str, Optional[Union[Num, Dict[str, Num]]]],
         weight: Num = None,
-        not_increment_count: bool = False,
     ) -> None:
         assert check_argument_types()
         if self._finished:
             raise RuntimeError("Already finished")
-        if not not_increment_count:
+        if len(self._seen_keys_in_the_step) == 0:
+            # Increment count as the first register in this step
             self.total_count += 1
             self.count += 1
 
         for key2, v in stats.items():
             if key2 in _reserved:
                 raise RuntimeError(f"{key2} is reserved.")
+            if key2 in self._seen_keys_in_the_step:
+                raise RuntimeError(f"{key2} is registered twice.")
             if v is None:
                 v = np.nan
             r = to_reported_value(v, weight)
-            self.stats[key2].append(r)
+
+            if key2 not in self.stats:
+                nan = to_reported_value(np.nan, None if weight is None else 0)
+                self.stats[key2].extend(
+                    r if i == self.count - 1 else nan for i in range(self.count)
+                )
+            else:
+                self.stats[key2].append(r)
+            self._seen_keys_in_the_step.add(key2)
 
     def log_message(self, start: int = None, end: int = None) -> str:
         if self._finished:
@@ -170,15 +196,11 @@ class SubReporter:
 
         message = f"{self.epoch}epoch:{self.key}:" f"{start + 1}-{end}batch: "
 
-        for idx, (key2, stats) in enumerate(self.stats.items()):
-            if len(stats) != self.count:
-                raise RuntimeError(
-                    "The registered stats must have same number of values: "
-                    f"{len(stats)} != {self.count}"
-                )
+        for idx, (key2, stats_list) in enumerate(self.stats.items()):
+            assert len(stats_list) == self.count, (len(stats_list), self.count)
             # values: List[ReportValue]
-            values = stats[start:end]
-            if idx != 0 and idx != len(stats):
+            values = stats_list[start:end]
+            if idx != 0 and idx != len(stats_list):
                 message += ", "
 
             v = aggregate(values)
@@ -196,14 +218,10 @@ class SubReporter:
         if start < 0:
             start = self.count + start
 
-        for key2, stats in self.stats.items():
-            if len(stats) != self.count:
-                raise RuntimeError(
-                    "The registered stats must have same number of values: "
-                    f"{len(stats)} != {self.count}"
-                )
+        for key2, stats_list in self.stats.items():
+            assert len(stats_list) == self.count, (len(stats_list), self.count)
             # values: List[ReportValue]
-            values = stats[start:]
+            values = stats_list[start:]
             v = aggregate(values)
             summary_writer.add_scalar(key2, v, self.total_count)
 
@@ -215,7 +233,7 @@ class SubReporter:
         start = time.perf_counter()
         yield start
         t = time.perf_counter() - start
-        self.register({name: t}, not_increment_count=True)
+        self.register({name: t})
 
     def measure_iter_time(self, iterable, name: str):
         iterator = iter(iterable)
@@ -224,7 +242,7 @@ class SubReporter:
                 start = time.perf_counter()
                 retval = next(iterator)
                 t = time.perf_counter() - start
-                self.register({name: t}, not_increment_count=True)
+                self.register({name: t})
                 yield retval
             except StopIteration:
                 break
diff --git a/espnet2/train/trainer.py b/espnet2/train/trainer.py
index 3b697aa36..da0d3ac84 100644
--- a/espnet2/train/trainer.py
+++ b/espnet2/train/trainer.py
@@ -378,7 +378,6 @@ class Trainer:
             batch = to_device(batch, "cuda" if ngpu > 0 else "cpu")
             if no_forward_run:
                 all_steps_are_invalid = False
-                reporter.register({})
                 continue
 
             with reporter.measure_time("forward_time"):
@@ -458,11 +457,11 @@ class Trainer:
                         },
                         train_time=time.perf_counter() - start_time,
                     ),
-                    # Suppress to increment the internal counter.
-                    not_increment_count=True,
                 )
                 start_time = time.perf_counter()
 
+            # NOTE(kamo): Call log_message() after increment_count()
+            reporter.next()
             if iiter % log_interval == 0:
                 logging.info(reporter.log_message(-log_interval))
                 if summary_writer is not None:
@@ -503,7 +502,6 @@ class Trainer:
 
             batch = to_device(batch, "cuda" if ngpu > 0 else "cpu")
             if no_forward_run:
-                reporter.register({})
                 continue
 
             _, stats, weight = model(**batch)
@@ -513,6 +511,7 @@ class Trainer:
                 stats, weight = recursive_average(stats, weight, distributed)
 
             reporter.register(stats, weight)
+            reporter.next()
 
         else:
             if distributed:
@@ -591,6 +590,4 @@ class Trainer:
                         summary_writer.add_figure(
                             f"{k}_{id_}", fig, reporter.get_epoch()
                         )
-
-                    # Dummy register() stimulates to increment the counter
-                    reporter.register({})
+            reporter.next()
diff --git a/test/espnet2/train/test_reporter.py b/test/espnet2/train/test_reporter.py
index 0774e2ac8..976024e4c 100644
--- a/test/espnet2/train/test_reporter.py
+++ b/test/espnet2/train/test_reporter.py
@@ -30,6 +30,8 @@ def test_register(weight1, weight2):
             "torch": torch.rand(1),
         }
         sub.register(stats1, weight1)
+        sub.next()
+
         stats2 = {
             "float": 0.3,
             "int": 100,
@@ -37,6 +39,7 @@ def test_register(weight1, weight2):
             "torch": torch.rand(1),
         }
         sub.register(stats2, weight2)
+        sub.next()
         assert sub.get_epoch() == 1
     with pytest.raises(RuntimeError):
         sub.register({})
@@ -69,6 +72,7 @@ def test_sort_epochs_and_values(mode):
         reporter.set_epoch(e + 1)
         with reporter.observe(key1) as sub:
             sub.register(stats_list[e])
+            sub.next()
     if mode not in ("min", "max"):
         with pytest.raises(ValueError):
             reporter.sort_epochs_and_values(key1, "aa", mode)
@@ -97,6 +101,7 @@ def test_sort_epochs_and_values_no_key():
         reporter.set_epoch(e + 1)
         with reporter.observe(key1) as sub:
             sub.register(stats_list[e])
+            sub.next()
     with pytest.raises(KeyError):
         reporter.sort_epochs_and_values("foo", "bar", "min")
 
@@ -116,6 +121,7 @@ def test_sort_values():
         reporter.set_epoch(e + 1)
         with reporter.observe(key1) as sub:
             sub.register(stats_list[e])
+            sub.next()
     sort_values = reporter.sort_values(key1, "aa", mode)
 
     desired = sorted([stats_list[e]["aa"] for e in range(len(stats_list))],)
@@ -133,6 +139,7 @@ def test_sort_epochs():
         reporter.set_epoch(e + 1)
         with reporter.observe(key1) as sub:
             sub.register(stats_list[e])
+            sub.next()
     sort_values = reporter.sort_epochs(key1, "aa", mode)
 
     desired = sorted(
@@ -153,6 +160,7 @@ def test_best_epoch():
         reporter.set_epoch(e + 1)
         with reporter.observe(key1) as sub:
             sub.register(stats_list[e])
+            sub.next()
     best_epoch = reporter.get_best_epoch(key1, "aa", mode)
     assert best_epoch == 3
 
@@ -169,6 +177,7 @@ def test_check_early_stopping():
         reporter.set_epoch(e + 1)
         with reporter.observe(key1) as sub:
             sub.register(stats_list[e])
+            sub.next()
         truefalse = reporter.check_early_stopping(patience, key1, "aa", mode)
         results.append(truefalse)
     assert results == [False, False, False, True]
@@ -188,8 +197,10 @@ def test_logging(tmp_path):
         reporter.set_epoch(e + 1)
         with reporter.observe(key1) as sub:
             sub.register(stats_list[e])
+            sub.next()
         with reporter.observe(key2) as sub:
             sub.register(stats_list[e])
+            sub.next()
             logging.info(sub.log_message())
             logging.info(sub.log_message(-1))
             logging.info(sub.log_message(0, 1))
@@ -201,11 +212,9 @@ def test_logging(tmp_path):
 
     with reporter.observe(key1) as sub:
         sub.register({"aa": 0.1, "bb": 0.4})
+        sub.next()
         sub.register({"aa": 0.1})
-        with pytest.raises(RuntimeError):
-            logging.info(sub.log_message())
-        with pytest.raises(RuntimeError):
-            sub.tensorboard_add_scalar(writer)
+        sub.next()
 
 
 def test_has_key():
@@ -215,6 +224,7 @@ def test_has_key():
     with reporter.observe(key1) as sub:
         stats1 = {"aa": 0.6}
         sub.register(stats1)
+        sub.next()
     assert reporter.has(key1, "aa")
 
 
@@ -225,6 +235,7 @@ def test_get_Keys():
     with reporter.observe(key1) as sub:
         stats1 = {"aa": 0.6}
         sub.register(stats1)
+        sub.next()
     assert reporter.get_keys() == (key1,)
 
 
@@ -235,6 +246,7 @@ def test_get_Keys2():
     with reporter.observe(key1) as sub:
         stats1 = {"aa": 0.6}
         sub.register(stats1)
+        sub.next()
     assert reporter.get_keys2(key1) == ("aa",)
 
 
@@ -245,16 +257,19 @@ def test_matplotlib_plot(tmp_path: Path):
     with reporter.observe(key1) as sub:
         stats1 = {"aa": 0.6}
         sub.register(stats1)
+        sub.next()
 
     reporter.set_epoch(1)
     with reporter.observe(key1) as sub:
         # Skip epoch=2
         sub.register({})
+        sub.next()
 
     reporter.set_epoch(3)
     with reporter.observe(key1) as sub:
         stats1 = {"aa": 0.6}
         sub.register(stats1)
+        sub.next()
 
     reporter.matplotlib_plot(tmp_path)
     assert (tmp_path / "aa.png").exists()
@@ -267,16 +282,19 @@ def test_tensorboard_add_scalar(tmp_path: Path):
     with reporter.observe(key1) as sub:
         stats1 = {"aa": 0.6}
         sub.register(stats1)
+        sub.next()
 
     reporter.set_epoch(1)
     with reporter.observe(key1) as sub:
         # Skip epoch=2
         sub.register({})
+        sub.next()
 
     reporter.set_epoch(3)
     with reporter.observe(key1) as sub:
         stats1 = {"aa": 0.6}
         sub.register(stats1)
+        sub.next()
 
     writer = SummaryWriter(tmp_path)
     reporter.tensorboard_add_scalar(writer)
@@ -288,9 +306,11 @@ def test_state_dict():
     with reporter.observe("train") as sub:
         stats1 = {"aa": 0.6}
         sub.register(stats1)
+        sub.next()
     with reporter.observe("eval") as sub:
         stats1 = {"bb": 0.6}
         sub.register(stats1)
+        sub.next()
     state = reporter.state_dict()
 
     reporter2 = Reporter()
@@ -310,9 +330,12 @@ def test_total_count():
     assert reporter.get_epoch() == 2
     with reporter.observe("train", 1) as sub:
         sub.register({})
+        sub.next()
     with reporter.observe("train", 2) as sub:
         sub.register({})
+        sub.next()
         sub.register({})
+        sub.next()
         assert sub.get_total_count() == 3
 
 
@@ -342,20 +365,24 @@ def test_register_array():
     with reporter.observe("train", 1) as sub:
         with pytest.raises(ValueError):
             sub.register({"a": np.array([0, 1])})
+            sub.next()
         with pytest.raises(ValueError):
-            sub.register({"a": 1}, weight=np.array([1, 2]))
+            sub.register({"b": 1}, weight=np.array([1, 2]))
+            sub.next()
 
 
 def test_zero_weight():
     reporter = Reporter()
     with reporter.observe("train", 1) as sub:
         sub.register({"a": 1}, weight=0)
+        sub.next()
 
 
 def test_register_nan():
     reporter = Reporter()
     with reporter.observe("train", 1) as sub:
         sub.register({"a": np.nan}, weight=1.0)
+        sub.next()
 
 
 def test_no_register():
@@ -368,8 +395,10 @@ def test_mismatch_key2():
     reporter = Reporter()
     with reporter.observe("train", 1) as sub:
         sub.register({"a": 2})
+        sub.next()
     with reporter.observe("train", 2) as sub:
         sub.register({"b": 3})
+        sub.next()
 
 
 def test_reserved():
@@ -377,8 +406,10 @@ def test_reserved():
     with reporter.observe("train", 1) as sub:
         with pytest.raises(RuntimeError):
             sub.register({"time": 2})
+            sub.next()
         with pytest.raises(RuntimeError):
             sub.register({"total_count": 3})
+            sub.next()
 
 
 def test_different_type():
@@ -386,13 +417,16 @@ def test_different_type():
     with pytest.raises(ValueError):
         with reporter.observe("train", 1) as sub:
             sub.register({"a": 2}, weight=1)
+            sub.next()
             sub.register({"a": 3})
+            sub.next()
 
 
 def test_start_middle_epoch():
     reporter = Reporter()
     with reporter.observe("train", 2) as sub:
         sub.register({"a": 3})
+        sub.next()
 
 
 def test__plot_stats_input_str():
@@ -420,10 +454,11 @@ def test_measure_time():
     with reporter.observe("train", 2) as sub:
         with sub.measure_time("foo"):
             pass
+        sub.next()
 
 
 def test_measure_iter_time():
     reporter = Reporter()
     with reporter.observe("train", 2) as sub:
         for _ in sub.measure_iter_time(range(3), "foo"):
-            pass
+            sub.next()