From 701adfa2a53093c7e8434e5c14875680b43a2777 Mon Sep 17 00:00:00 2001
From: RedRayz <71994877+RedRayz@users.noreply.github.com>
Date: Fri, 5 Jan 2024 21:32:00 +0900
Subject: [PATCH 1/2] Show total progress and don't evaluate when script starts
 training

---
 train_ms.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/train_ms.py b/train_ms.py
index a1e72d7..34be92b 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -327,7 +327,7 @@ def run():
             utils.get_steps(utils.latest_checkpoint_path(model_dir, "G_*.pth"))
         )
         logger.info(
-            f"******************检测到模型存在，epoch为 {epoch_str}，gloabl step为 {global_step}*********************"
+            f"******************Found the model. Current epoch is {epoch_str}, gloabl step is {global_step}*********************"
         )
     else:
         try:
@@ -365,6 +365,10 @@ def run():
         scheduler_dur_disc = None
     scaler = GradScaler(enabled=hps.train.bf16_run)
 
+    diff = abs(epoch_str * len(train_loader) - (hps.train.epochs + 1) * len(train_loader))
+    pbar = tqdm(total=global_step + diff, initial=global_step, smoothing=0.05, file=SAFE_STDOUT)
+    initial_step = global_step
+
     for epoch in range(epoch_str, hps.train.epochs + 1):
         if rank == 0:
             train_and_evaluate(
@@ -379,6 +383,8 @@ def run():
                 [train_loader, eval_loader],
                 logger,
                 [writer, writer_eval],
+                pbar,
+                initial_step,
             )
         else:
             train_and_evaluate(
@@ -393,6 +399,8 @@ def run():
                 [train_loader, None],
                 None,
                 None,
+                pbar,
+                initial_step,
             )
         scheduler_g.step()
         scheduler_d.step()
@@ -433,6 +441,8 @@ def run():
                 for_infer=True,
             )
 
+    pbar.close()
+
 
 def train_and_evaluate(
     rank,
@@ -446,6 +456,8 @@ def train_and_evaluate(
     loaders,
     logger,
     writers,
+    pbar: tqdm,
+    initial_step: int
 ):
     net_g, net_d, net_dur_disc = nets
     optim_g, optim_d, optim_dur_disc = optims
@@ -475,7 +487,7 @@ def train_and_evaluate(
         ja_bert,
         en_bert,
         style_vec,
-    ) in enumerate(tqdm(train_loader, file=SAFE_STDOUT)):
+    ) in enumerate(train_loader):
         if net_g.module.use_noise_scaled_mas:
             current_mas_noise_scale = (
                 net_g.module.mas_noise_scale_initial
@@ -663,7 +675,7 @@ def train_and_evaluate(
                     scalars=scalar_dict,
                 )
 
-            if global_step % hps.train.eval_interval == 0 and global_step != 0:
+            if global_step % hps.train.eval_interval == 0 and global_step != 0 and initial_step != global_step:
                 evaluate(hps, net_g, eval_loader, writer_eval)
                 utils.save_checkpoint(
                     net_g,
@@ -706,11 +718,12 @@ def train_and_evaluate(
                 )
 
         global_step += 1
+        if pbar is not None:
+            pbar.set_description("Epoch {}({:.0f}%)/{}".format(epoch, 100.0 * batch_idx / len(train_loader), hps.train.epochs))
+            pbar.update()
     # 本家ではこれをスピードアップのために消すと書かれていたので、一応消してみる
     # gc.collect()
     # torch.cuda.empty_cache()
-    if rank == 0:
-        logger.info(f"====> Epoch: {epoch}, step: {global_step}")
 
 
 def evaluate(hps, generator, eval_loader, writer_eval):

From bf9a0a12e75a66a3336ffe61e39294dd0bbcb442 Mon Sep 17 00:00:00 2001
From: RedRayz <71994877+RedRayz@users.noreply.github.com>
Date: Sat, 6 Jan 2024 05:22:16 +0900
Subject: [PATCH 2/2] Add --no_progress_bar option

---
 train_ms.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/train_ms.py b/train_ms.py
index 34be92b..cbf6d72 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -107,6 +107,11 @@ def run():
         action="store_true",
         help="Skip saving default style config and mean vector.",
     )
+    parser.add_argument(
+        '--no_progress_bar',
+        action='store_true',
+        help='Do not show the progress bar while training.'
+    )
     args = parser.parse_args()
     model_dir = os.path.join(args.model, config.train_ms_config.model_dir)
     if not os.path.exists(model_dir):
@@ -364,9 +369,12 @@ def run():
     else:
         scheduler_dur_disc = None
     scaler = GradScaler(enabled=hps.train.bf16_run)
+    logger.info("Start training.")
 
     diff = abs(epoch_str * len(train_loader) - (hps.train.epochs + 1) * len(train_loader))
-    pbar = tqdm(total=global_step + diff, initial=global_step, smoothing=0.05, file=SAFE_STDOUT)
+    pbar = None
+    if not args.no_progress_bar:
+        pbar = tqdm(total=global_step + diff, initial=global_step, smoothing=0.05, file=SAFE_STDOUT)
     initial_step = global_step
 
     for epoch in range(epoch_str, hps.train.epochs + 1):
@@ -441,7 +449,8 @@ def run():
                 for_infer=True,
             )
 
-    pbar.close()
+    if pbar is not None:
+        pbar.close()
 
 
 def train_and_evaluate(
@@ -724,7 +733,8 @@ def train_and_evaluate(
     # 本家ではこれをスピードアップのために消すと書かれていたので、一応消してみる
     # gc.collect()
     # torch.cuda.empty_cache()
-
+    if pbar is None and rank == 0:
+        logger.info(f"====> Epoch: {epoch}, step: {global_step}")
 
 def evaluate(hps, generator, eval_loader, writer_eval):
     generator.eval()