Restore gc and emptycache when training

train_ms.py

@@ -1,5 +1,6 @@
 import argparse
 import datetime
+import gc
 import os
 import platform
 
@@ -757,8 +758,9 @@ def train_and_evaluate(
             )
             pbar.update()
     # 本家ではこれをスピードアップのために消すと書かれていたので、一応消してみる
-    # gc.collect()
+    # と思ったけどメモリ使用量が減るかもしれないのでつけてみる
-    # torch.cuda.empty_cache()
+    gc.collect()
+    torch.cuda.empty_cache()
     if pbar is None and rank == 0:
         logger.info(f"====> Epoch: {epoch}, step: {global_step}")
 

train_ms_jp_extra.py

@@ -1,5 +1,6 @@
 import argparse
 import datetime
+import gc
 import os
 import platform
 
@@ -589,7 +590,7 @@ def train_and_evaluate(
         language,
         bert,
         style_vec,
-    ) in enumerate(tqdm(train_loader)):
+    ) in enumerate(train_loader):
         if net_g.module.use_noise_scaled_mas:
             current_mas_noise_scale = (
                 net_g.module.mas_noise_scale_initial
@@ -899,8 +900,9 @@ def train_and_evaluate(
             )
             pbar.update()
     # 本家ではこれをスピードアップのために消すと書かれていたので、一応消してみる
-    # gc.collect()
+    # と思ったけどメモリ使用量が減るかもしれないのでつけてみる
-    # torch.cuda.empty_cache()
+    gc.collect()
+    torch.cuda.empty_cache()
     if pbar is None and rank == 0:
         logger.info(f"====> Epoch: {epoch}, step: {global_step}")
 

webui_train.py

@@ -302,7 +302,7 @@ def train(model_name, skip_style=False, use_jp_extra=True):
     cmd = [train_py, "--config", config_path, "--model", dataset_path]
     if skip_style:
         cmd.append("--skip_default_style")
-    success, message = run_script_with_log(cmd)
+    success, message = run_script_with_log(cmd, ignore_warning=True)
     if not success:
         logger.error(f"Train failed.")
         return False, f"Error: 学習に失敗しました:\n{message}"