Merge pull request #1171 from MouseLand/allow_save_each

allow save_each when training model

Author: mrariden

cellpose/__main__.py

@@ -111,6 +111,9 @@ def main():
     else:
         normalize = (not args.no_norm)
 
+    if args.save_each:
+        if not args.save_every:
+            raise ValueError("ERROR: --save_each requires --save_every")
 
     if len(args.image_path) > 0 and args.train:
         raise ValueError("ERROR: cannot train model with single image input")
@@ -163,7 +166,9 @@ def _train_cellposemodel_cli(args, logger, image_filter, device, pretrained_mode
             min_train_masks=args.min_train_masks,
             nimg_per_epoch=args.nimg_per_epoch,
             nimg_test_per_epoch=args.nimg_test_per_epoch,
-            save_path=os.path.realpath(args.dir), save_every=args.save_every,
+            save_path=os.path.realpath(args.dir), 
+            save_every=args.save_every,
+            save_each=args.save_each,
             model_name=args.model_name_out)[0]
     model.pretrained_model = cpmodel_path
     logger.info(">>>> model trained and saved to %s" % cpmodel_path)

cellpose/cli.py

@@ -222,6 +222,9 @@ def get_arg_parser():
     training_args.add_argument(
         "--save_every", default=100, type=int,
         help="number of epochs to skip between saves. Default: %(default)s")
+    training_args.add_argument(
+        "--save_each", action="store_true",
+        help="wether or not to save each epoch. Must also use --save_every. (default: False)")
     training_args.add_argument(
         "--model_name_out", default=None, type=str,
         help="Name of model to save as, defaults to name describing model architecture. "