fix the bug of stop training around 4 hours.

Signed-off-by: lawrence-cj <[email protected]>

diffusion/utils/config.py

@@ -141,6 +141,7 @@ class TrainingConfig(BaseConfig):
     load_mask_index: bool = False
     snr_loss: bool = False
     real_prompt_ratio: float = 1.0
+    training_hours: float = 10000.0
     save_image_epochs: int = 1
     save_model_epochs: int = 1
     save_model_steps: int = 1000000

train_scripts/train.py

@@ -448,7 +448,10 @@ def train(config, args, accelerator, model, optimizer, lr_scheduler, train_datal
 
             if loss_nan_timer > 20:
                 raise ValueError("Loss is NaN too much times. Break here.")
-            if global_step % config.train.save_model_steps == 0 or (time.time() - training_start_time) / 3600 > 3.8:
+            if (
+                global_step % config.train.save_model_steps == 0
+                or (time.time() - training_start_time) / 3600 > config.train.training_hours
+            ):
                 accelerator.wait_for_everyone()
                 if accelerator.is_main_process:
                     os.umask(0o000)
@@ -469,7 +472,7 @@ def train(config, args, accelerator, model, optimizer, lr_scheduler, train_datal
                             f.write(osp.join(config.work_dir, "config.py") + "\n")
                             f.write(ckpt_saved_path)
 
-                if (time.time() - training_start_time) / 3600 > 3.8:
+                if (time.time() - training_start_time) / 3600 > config.train.training_hours:
                     logger.info(f"Stopping training at epoch {epoch}, step {global_step} due to time limit.")
                     return
             if config.train.visualize and (global_step % config.train.eval_sampling_steps == 0 or (step + 1) == 1):