Do not use iterations done in json log

Author: dmitriy-serdyuk

blocks/log/json.py

@@ -65,46 +65,55 @@ class JSONLinesLog(TrainingLogBase):
     def __init__(self, filename='log.jsonl.gz', maxlen=101, formatter=None,
                  **kwargs):
         self.status = {}
-        TrainingLogBase.__init__(self)
+        super(JSONLinesLog, self).__init__()
         if os.path.isfile(filename):
             os.remove(filename)
         self.logger = PicklableLogger(
             filename=filename, maxlen=maxlen, formatter=formatter, **kwargs)
         self.local_cache = deque()
 
     def flush(self, iterations_done):
+        if iterations_done < 0:
+            raise ValueError
         if len(self.local_cache) > 0:
             self.logger.log({'iterations_done': iterations_done,
                              'reports': self.local_cache.popleft()})
 
     def __getitem__(self, time):
         self._check_time(time)
-        iterations_done = self.status.get('iterations_done', -1)
+        logger_len = self.inner_logger_len()
+        total_length = logger_len + len(self.local_cache)
 
         # Flush local cache
         while len(self.local_cache) > 1:
-            self.flush(iterations_done - len(self.local_cache) + 1)
+            self.flush(total_length - len(self.local_cache))
+        logger_len = self.inner_logger_len()
 
-        total_length = len(self.logger) + len(self.local_cache)
         if time >= total_length:
             # Need to create new item in local cache
             self.local_cache.extend(
                 [{} for _ in range(time - total_length + 1)])
-        last_logged_element = len(self.logger)
-        if time < last_logged_element:
+        if time < logger_len:
             try:
-                assert self.logger[time]['iterations_done'] == time
+                if not self.logger[time]['iterations_done'] == time:
+                    raise ValueError('iterations done')
                 return self.logger[time]['reports']
             except IndexError:
                 raise ValueError(
                     'cannot get past log entries for JSON log, max log length '
                     'in memory is: {}'.format(
                         self.logger.logger_kwargs['maxlen']))
-        if time >= last_logged_element:
-            return self.local_cache[time - last_logged_element]
+        if time >= logger_len:
+            return self.local_cache[time - logger_len]
+
+    def inner_logger_len(self):
+        try:
+            return len(self.logger)
+        except AttributeError:
+            return 0
 
     def __len__(self):
-        return len(self.logger) + len(self.local_cache)
+        return self.inner_logger_len() + len(self.local_cache)
 
     def __setitem__(self, time, value):
         raise ValueError('cannot manually change JSON Lines log')
@@ -115,3 +124,6 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.flush(self.status.get('iterations_done', -1))
         self.logger.close()
+
+    def __iter__(self):
+        return iter([self[i] for i in range(len(self))])

tests/extensions/test_training.py

@@ -90,33 +90,34 @@ def test_track_the_best():
     extension = TrackTheBest("cost")
     extension.main_loop = main_loop
 
-    main_loop.status['epochs_done'] += 1
-    main_loop.status['iterations_done'] += 10
-    main_loop.log.current_row['cost'] = 5
-    extension.dispatch('after_epoch')
-    assert main_loop.status['best_cost'] == 5
-    assert main_loop.log.current_row['cost_best_so_far']
-
-    main_loop.status['epochs_done'] += 1
-    main_loop.status['iterations_done'] += 10
-    main_loop.log.current_row['cost'] = 6
-    extension.dispatch('after_epoch')
-    assert main_loop.status['best_cost'] == 5
-    assert main_loop.log.current_row.get('cost_best_so_far', None) is None
-
-    main_loop.status['epochs_done'] += 1
-    main_loop.status['iterations_done'] += 10
-    main_loop.log.current_row['cost'] = 5
-    extension.dispatch('after_epoch')
-    assert main_loop.status['best_cost'] == 5
-    assert main_loop.log.current_row.get('cost_best_so_far', None) is None
-
-    main_loop.status['epochs_done'] += 1
-    main_loop.status['iterations_done'] += 10
-    main_loop.log.current_row['cost'] = 4
-    extension.dispatch('after_epoch')
-    assert main_loop.status['best_cost'] == 4
-    assert main_loop.log.current_row['cost_best_so_far']
+    with main_loop.log:
+        main_loop.status['epochs_done'] += 1
+        main_loop.status['iterations_done'] += 10
+        main_loop.log.current_row['cost'] = 5
+        extension.dispatch('after_epoch')
+        assert main_loop.status['best_cost'] == 5
+        assert main_loop.log.current_row['cost_best_so_far']
+
+        main_loop.status['epochs_done'] += 1
+        main_loop.status['iterations_done'] += 10
+        main_loop.log.current_row['cost'] = 6
+        extension.dispatch('after_epoch')
+        assert main_loop.status['best_cost'] == 5
+        assert main_loop.log.current_row.get('cost_best_so_far', None) is None
+
+        main_loop.status['epochs_done'] += 1
+        main_loop.status['iterations_done'] += 10
+        main_loop.log.current_row['cost'] = 5
+        extension.dispatch('after_epoch')
+        assert main_loop.status['best_cost'] == 5
+        assert main_loop.log.current_row.get('cost_best_so_far', None) is None
+
+        main_loop.status['epochs_done'] += 1
+        main_loop.status['iterations_done'] += 10
+        main_loop.log.current_row['cost'] = 4
+        extension.dispatch('after_epoch')
+        assert main_loop.status['best_cost'] == 4
+        assert main_loop.log.current_row['cost_best_so_far']
 
 
 class WriteCostExtension(TrainingExtension):

tests/test_log.py

@@ -3,30 +3,43 @@
 
 from numpy.testing import assert_raises
 
-from blocks.log import TrainingLog
+from blocks.log import TrainingLog, JSONLinesLog
 from blocks.serialization import load, dump
 
 
-def test_training_log():
-    log = TrainingLog()
+def run_log(log):
+    with log:
+        # test basic writing capabilities
+        log[0]['field'] = 45
+        assert log[0]['field'] == 45
+        assert log[1] == {}
+        assert log.current_row['field'] == 45
+        log.status['iterations_done'] += 1
+        assert log.status['iterations_done'] == 1
+        assert log.previous_row['field'] == 45
+
+        assert_raises(ValueError, getitem, log, -1)
+
+        # test iteration
+        assert len(list(log)) == 2
+
 
-    # test basic writing capabilities
-    log[0]['field'] = 45
-    assert log[0]['field'] == 45
-    assert log[1] == {}
-    assert log.current_row['field'] == 45
-    log.status['iterations_done'] += 1
-    assert log.status['iterations_done'] == 1
-    assert log.previous_row['field'] == 45
+def test_json_lines_log():
+    log = JSONLinesLog(maxlen=2)
+    run_log(log)
 
-    assert_raises(ValueError, getitem, log, -1)
 
-    # test iteration
-    assert len(list(log)) == 2
+def test_training_log():
+    log = TrainingLog()
+    run_log(log)
 
 
 def test_pickle_log():
-    log1 = TrainingLog()
+    log = TrainingLog()
+    pickle_log(log)
+
+
+def pickle_log(log1):
     with open('log1.tar', 'wb') as f:
         dump(log1, f)
     with open('log1.tar', 'rb') as f: