Split from PR#968. add --keep-best-checkpoints (#990)

Summary:
Fixes fairinternal/fairseq-py#968. Split --keep-best-checkpoints from the original request.
Use scores as the names to save the checkpoints
Pull Request resolved: fairinternal/fairseq-py#990

Differential Revision: D19411250

Pulled By: MultiPath

fbshipit-source-id: 82b0db614208eee54c9c0e470ad7faa6481747d5

Author: MultiPath

fairseq/checkpoint_utils.py

@@ -53,6 +53,15 @@ def is_better(a, b):
         not hasattr(save_checkpoint, "best")
         or is_better(val_loss, save_checkpoint.best)
     )
+    checkpoint_conds["checkpoint.best_{}_{:.2f}.pt".format(
+        args.best_checkpoint_metric, val_loss)] = (
+        val_loss is not None
+        and args.keep_best_checkpoints > 0
+        and (
+            not hasattr(save_checkpoint, "best")
+            or is_better(val_loss, save_checkpoint.best)
+        )
+    )
     checkpoint_conds["checkpoint_last.pt"] = not args.no_last_checkpoints
 
     extra_state = {"train_iterator": epoch_itr.state_dict(), "val_loss": val_loss}
@@ -69,8 +78,8 @@ def is_better(a, b):
 
         write_timer.stop()
         print(
-            "| saved checkpoint {} (epoch {} @ {} updates) (writing took {} seconds)".format(
-                checkpoints[0], epoch, updates, write_timer.sum
+            "| saved checkpoint {} (epoch {} @ {} updates, score {}) (writing took {} seconds)".format(
+                checkpoints[0], epoch, updates, val_loss, write_timer.sum
             )
         )
 
@@ -90,6 +99,15 @@ def is_better(a, b):
             if os.path.lexists(old_chk):
                 os.remove(old_chk)
 
+    if args.keep_best_checkpoints > 0:
+        # only keep the best N checkpoints according to validation metric
+        checkpoints = checkpoint_paths(
+            args.save_dir, pattern=r"checkpoint\.best_{}_(\d+\.?\d*)\.pt".format(args.best_checkpoint_metric))
+        if not args.maximize_best_checkpoint_metric:
+            checkpoints = checkpoints[::-1]
+        for old_chk in checkpoints[args.keep_best_checkpoints:]:
+            if os.path.lexists(old_chk):
+                os.remove(old_chk)
 
 def load_checkpoint(args, trainer, **passthrough_args):
     """
@@ -202,7 +220,7 @@ def checkpoint_paths(path, pattern=r"checkpoint(\d+)\.pt"):
     for i, f in enumerate(files):
         m = pt_regexp.fullmatch(f)
         if m is not None:
-            idx = int(m.group(1)) if len(m.groups()) > 0 else i
+            idx = float(m.group(1)) if len(m.groups()) > 0 else i
             entries.append((idx, m.group(0)))
     return [os.path.join(path, x[1]) for x in sorted(entries, reverse=True)]
 

fairseq/options.py

@@ -403,6 +403,8 @@ def add_checkpoint_args(parser):
                        help='keep the last N checkpoints saved with --save-interval-updates')
     group.add_argument('--keep-last-epochs', type=int, default=-1, metavar='N',
                        help='keep last N epoch checkpoints')
+    group.add_argument('--keep-best-checkpoints', type=int, default=-1, metavar='N',
+                       help='keep best N checkpoints based on scores')
     group.add_argument('--no-save', action='store_true',
                        help='don\'t save models or checkpoints')
     group.add_argument('--no-epoch-checkpoints', action='store_true',