From 4389ae75c885d78c5f989cb4ce5c92eeebe9afeb Mon Sep 17 00:00:00 2001 From: xander Date: Tue, 14 May 2024 15:21:48 +0200 Subject: [PATCH] force quit training after max steps --- main.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/main.py b/main.py index abedb0a..f852dc2 100755 --- a/main.py +++ b/main.py @@ -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}"