force quit training after max steps

2024-05-14 15:21:48 +02:00
parent d66863b4cb
commit 4389ae75c8
1 changed files with 6 additions and 4 deletions
@@ -430,14 +430,16 @@ def train(
            
            images_done += config.train_batch_size
            global_step += 1
-
-            if global_step % 100 == 0:
-                print(f" ---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r")
-
+            
            if global_step % (config.max_train_steps//20) == 0:
                progress = (global_step / config.max_train_steps) + 0.05
+                print(f" ---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r")
                yield np.min((progress, 1.0))

+            if global_step > config.max_train_steps:
+                print("Reached max steps, stopping training!")
+                break
+
    # final_save
    if (global_step - last_save_step) > 26:
        output_save_dir = f"{checkpoint_dir}/checkpoint-{global_step}"