Repository: awslabs/sockeye

Issue:
Release 2.1.16: metrics file loses a line every time training is stopped and then continued
Let's say you train for 10 checkpoints and stop. The metrics file will have info for 10 checkpoints, and the latest params file will be params.00010.

Now, suppose you continue training and stop after an additional 10 checkpoints. The metrics file will have info for 19 checkpoints, and the latest params file will be params.00020. The line that was lost was the line corresponding to params.00010. All of the lines from the 10'th on are now incorrect. Line (9+n) of the metrics file actually contains information about checkpoint (9 + n + 1).


Discussion:

Resolving patch:
diff --git a/sockeye/__init__.py b/sockeye/__init__.py
--- a/sockeye/__init__.py
+++ b/sockeye/__init__.py
@@ -11,4 +11,4 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-__version__ = '2.1.19'
+__version__ = '2.1.20'
diff --git a/sockeye/training.py b/sockeye/training.py
--- a/sockeye/training.py
+++ b/sockeye/training.py
@@ -96,9 +96,9 @@ class TrainState:
 
     _pickle_slots = ['num_not_improved', 'epoch', 'checkpoint', 'best_checkpoint', 'batches',
                      'updates', 'samples', 'gradient_norm', 'metrics', 'start_tic', '_tic_last_time_elapsed',
-                     '_time_elapsed', 'early_stopping_metric', 'best_metric', 'best_metric_history', 
+                     '_time_elapsed', 'early_stopping_metric', 'best_metric', 'best_metric_history',
                      'best_checkpoint', 'converged', 'diverged']
- 
+
     __slots__ = _pickle_slots + ['gradients']
 
     def __init__(self, early_stopping_metric: str) -> None:
@@ -129,6 +129,7 @@ def save(self, fname: str):
         Saves this training state to fname.
         """
         self.update_time_elapsed()
+        assert len(self.metrics) == self.checkpoint
         with open(fname, "wb") as fp:
             pickle.dump(self, fp)
 
@@ -140,6 +141,7 @@ def load(fname: str) -> 'TrainState':
         with open(fname, "rb") as fp:
             state = pickle.load(fp)
             state._tic_last_time_elapsed = time.time()
+            assert len(state.metrics) == state.checkpoint
             return state
 
     def update_time_elapsed(self):
@@ -272,12 +274,13 @@ def fit(self,
                 if has_improved:
                     self._update_best_params()
                     self._save_trainer_states(self.best_optimizer_states_fname)
-                self._save_training_state(train_iter)
 
                 self._write_and_log_metrics(train_metrics=train_metrics, val_metrics=val_metrics)
                 for metric in train_metrics:
                     metric.reset()
 
+                self._save_training_state(train_iter)
+
                 if self.checkpoint_callback:
                     self.checkpoint_callback(self.state.checkpoint)
 
@@ -661,6 +664,11 @@ def _load_training_state(self, train_iter: data_io.BaseParallelSampleIter):
                  self.trainer._amp_loss_scaler._next_loss_scale,
                  self.trainer._amp_loss_scaler._unskipped) = pickle.load(fp)
 
+        logger.info("Training State: epoch=%d, checkpoint=%d batches=%d updates=%d best_metric=%.2f, " \
+                    "best_checkpoint=%d time_elapsed=%d" % (
+                        self.state.epoch, self.state.checkpoint, self.state.batches, self.state.updates,
+                        self.state.best_metric, self.state.best_checkpoint, self.state.time_elapsed))
+
     def _cleanup(self, keep_training_state=False):
         """
         Cleans parameter files, training state directory and waits for remaining decoding processes.


Test patch: