From 0f4c9b2732fc852adb3bdec8dbe72d36de6f9d60 Mon Sep 17 00:00:00 2001 From: Ron Date: Sun, 4 Jul 2021 09:59:13 +0300 Subject: [PATCH] fix number of iterations --- train_first_stage.py | 5 ++++- train_second_stage.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/train_first_stage.py b/train_first_stage.py index e90113d..981ff60 100755 --- a/train_first_stage.py +++ b/train_first_stage.py @@ -627,6 +627,8 @@ def train(args): print("Checkpoint saved") iter_cnt += 1 + if iter_cnt > args.iters: + break print("Training is done") save_file = os.path.join(args.out, 'checkpoint') @@ -648,7 +650,8 @@ def train(args): parser.add_argument('--g_lr', type=float, default=0.0001) parser.add_argument('--d_lr', type=float, default=0.0001) parser.add_argument('--bs', type=int, default=8) - parser.add_argument('--epoch', type=int, default=45000) + parser.add_argument('--epoch', type=int, default=2000) + parser.add_argument('--iters', type=int, default=45001) parser.add_argument('--resize_w', type=int, default=256) parser.add_argument('--resize_h', type=int, default=256) parser.add_argument('--num_kp', type=int, default=10) diff --git a/train_second_stage.py b/train_second_stage.py index 7628e0b..f3497c2 100755 --- a/train_second_stage.py +++ b/train_second_stage.py @@ -448,6 +448,8 @@ def train(args, opt=None): print("Checkpoint saved") iter_cnt += 1 + if iter_cnt > args.iters: + break print("Training is done") @@ -470,7 +472,8 @@ def train(args, opt=None): parser.add_argument('--g_lr', type=float, default=0.0001) parser.add_argument('--d_lr', type=float, default=0.0001) parser.add_argument('--bs', type=int, default=4) - parser.add_argument('--epoch', type=int, default=1000000) + parser.add_argument('--epoch', type=int, default=2000) + parser.add_argument('--iters', type=int, default=30001) parser.add_argument('--resize_w', type=int, default=256) parser.add_argument('--resize_h', type=int, default=256) parser.add_argument('--num_kp', type=int, default=10)