Improved checkpointing.

rmcpantoja · Dec 11, 2023 · bac6b8d · bac6b8d · Sobsz · Dec 11, 2023
1 parent e8edb3b
commit bac6b8d
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/src/python/piper_train/__main__.py b/src/python/piper_train/__main__.py
@@ -37,6 +37,12 @@ def main():
     Trainer.add_argparse_args(parser)
     VitsModel.add_model_specific_args(parser)
     parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument(
+        "--num_ckpt", default=1, help="# of ckpts saved."
+    )
+    parser.add_argument(
+        "--save_last", default=False, help="Always save the last checkpoint."
+    )
     args = parser.parse_args()
     _LOGGER.debug(args)
 
@@ -59,10 +65,18 @@ def main():
 
     trainer = Trainer.from_argparse_args(args)
     if args.checkpoint_epochs is not None:
-        trainer.callbacks = [ModelCheckpoint(every_n_epochs=args.checkpoint_epochs)]
+        trainer.callbacks = [ModelCheckpoint(
+            every_n_epochs=args.checkpoint_epochs,
+            filename="model",
+            save_top_k=args.num_ckpt,
+            save_last=args.save_last
+        )]
         _LOGGER.debug(
             "Checkpoints will be saved every %s epoch(s)", args.checkpoint_epochs
         )
+        _LOGGER.debug(
+            "%s Checkpoints will be saved", args.num_ckpt
+        )
 
     dict_args = vars(args)
     if args.quality == "x-low":