fix number of iterations

rmokady · Jul 4, 2021 · 0f4c9b2 · 0f4c9b2
1 parent 4eb3533
commit 0f4c9b2
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git 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
@@ -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)