From f14ce6e313e1bf785ea64923f94ede7ff716e951 Mon Sep 17 00:00:00 2001 From: Brigham Lysenko Date: Sun, 11 Aug 2019 15:55:14 -0600 Subject: [PATCH] Added pingpong option --- main.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/main.py b/main.py index 8093c77..3aef023 100644 --- a/main.py +++ b/main.py @@ -134,16 +134,23 @@ if __name__ == "__main__": Trainer.main(args, device_args) p = subparsers.add_parser( "train", help="Trainer") - p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of extracted SRC faceset.") - p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of extracted DST faceset.") - p.add_argument('--pretraining-data-dir', action=fixPathAction, dest="pretraining_data_dir", default=None, help="Optional dir of extracted faceset that will be used in pretraining mode.") + p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", + help="Dir of extracted SRC faceset.") + p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", + help="Dir of extracted DST faceset.") + p.add_argument('--pretraining-data-dir', action=fixPathAction, dest="pretraining_data_dir", default=None, + help="Optional dir of extracted faceset that will be used in pretraining mode.") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") - p.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") - p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.") + p.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), + help="Type of model") + p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, + help="Disable preview window.") p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+') + p.add_argument('--pingpong', dest="ping_pong", default=False, + help="Cycle between a batch size of 1 and the chosen batch size") p.set_defaults (func=process_train) def process_convert(arguments):