Merge pull request #1 from Raghava14/best_checkpoint

Added saving of best checkpoint

Author: sugeeth14

onmt/models/model_saver.py

@@ -37,7 +37,7 @@ def __init__(self, base_path, model, model_opt, fields, optim,
         if keep_checkpoint > 0:
             self.checkpoint_queue = deque([], maxlen=keep_checkpoint)
 
-    def save(self, step, moving_average=None):
+    def save(self, step, is_best, moving_average=None):
         """Main entry point for model saver
 
         It wraps the `_save` method with checks and apply `keep_checkpoint`
@@ -54,7 +54,7 @@ def save(self, step, moving_average=None):
                 model_params_data.append(param.data)
                 param.data = avg.data
 
-        chkpt, chkpt_name = self._save(step, save_model)
+        chkpt, chkpt_name = self._save(step, is_best, save_model)
         self.last_saved_step = step
 
         if moving_average:
@@ -97,7 +97,7 @@ def _rm_checkpoint(self, name):
 class ModelSaver(ModelSaverBase):
     """Simple model saver to filesystem"""
 
-    def _save(self, step, model):
+    def _save(self, step, is_best, model):
         model_state_dict = model.state_dict()
         model_state_dict = {k: v for k, v in model_state_dict.items()
                             if 'generator' not in k}
@@ -128,6 +128,10 @@ def _save(self, step, model):
         logger.info("Saving checkpoint %s_step_%d.pt" % (self.base_path, step))
         checkpoint_path = '%s_step_%d.pt' % (self.base_path, step)
         torch.save(checkpoint, checkpoint_path)
+        if is_best:
+            logger.info("Obtained best checkpoint, saving...")
+            best_checkpoint_path = '%s_best.pt' % (self.base_path)
+            torch.save(checkpoint, best_checkpoint_path)
         return checkpoint, checkpoint_path
 
     def _rm_checkpoint(self, name):

onmt/trainer.py

@@ -267,6 +267,17 @@ def train(self,
             if self.average_decay > 0 and i % self.average_every == 0:
                 self._update_average(step)
 
+            perplexity = report_stats.ppl()
+            if i == 0:
+                best_perplexity = perplexity
+                is_best = True
+            else:
+                if perplexity < best_perplexity:
+                    best_perplexity = perplexity
+                    is_best = True
+                else:
+                    is_best = False
+
             report_stats = self._maybe_report_training(
                 step, train_steps,
                 self.optim.learning_rate(),
@@ -297,13 +308,15 @@ def train(self,
             if (self.model_saver is not None
                 and (save_checkpoint_steps != 0
                      and step % save_checkpoint_steps == 0)):
-                self.model_saver.save(step, moving_average=self.moving_average)
+                self.model_saver.save(step, is_best,
+                                      moving_average=self.moving_average)
 
             if train_steps > 0 and step >= train_steps:
                 break
 
         if self.model_saver is not None:
-            self.model_saver.save(step, moving_average=self.moving_average)
+            self.model_saver.save(step, is_best,
+                                  moving_average=self.moving_average)
         return total_stats
 
     def validate(self, valid_iter, moving_average=None):