keep only last ckpt + black + fixes multi gpu tqdm

Author: Dschogo

train/utils.py

@@ -31,10 +31,7 @@ def go(model, bkey):
             try:
                 new_state_dict[k] = saved_state_dict[k]
                 if saved_state_dict[k].shape != state_dict[k].shape:
-                    print(
-                        "shape-%s-mismatch|need-%s|get-%s"
-                        % (k, state_dict[k].shape, saved_state_dict[k].shape)
-                    )  #
+                    print("shape-%s-mismatch|need-%s|get-%s" % (k, state_dict[k].shape, saved_state_dict[k].shape))  #
                     raise KeyError
             except:
                 # logger.info(traceback.format_exc())
@@ -52,9 +49,7 @@ def go(model, bkey):
 
     iteration = checkpoint_dict["iteration"]
     learning_rate = checkpoint_dict["learning_rate"]
-    if (
-        optimizer is not None and load_opt == 1
-    ):  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
+    if optimizer is not None and load_opt == 1:  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
         #   try:
         optimizer.load_state_dict(checkpoint_dict["optimizer"])
     #   except:
@@ -106,10 +101,7 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
         try:
             new_state_dict[k] = saved_state_dict[k]
             if saved_state_dict[k].shape != state_dict[k].shape:
-                print(
-                    "shape-%s-mismatch|need-%s|get-%s"
-                    % (k, state_dict[k].shape, saved_state_dict[k].shape)
-                )  #
+                print("shape-%s-mismatch|need-%s|get-%s" % (k, state_dict[k].shape, saved_state_dict[k].shape))  #
                 raise KeyError
         except:
             # logger.info(traceback.format_exc())
@@ -123,9 +115,7 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
 
     iteration = checkpoint_dict["iteration"]
     learning_rate = checkpoint_dict["learning_rate"]
-    if (
-        optimizer is not None and load_opt == 1
-    ):  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
+    if optimizer is not None and load_opt == 1:  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
         #   try:
         optimizer.load_state_dict(checkpoint_dict["optimizer"])
     #   except:
@@ -134,33 +124,39 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
     return model, optimizer, learning_rate, iteration
 
 
-def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
-    logger.info(
-        "Saving model and optimizer state at epoch {} to {}".format(
-            iteration, checkpoint_path
-        )
-    )
+def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path, checkpoint_type, delete_old=False):
+    # logger.info(
+    #     "Saving model and optimizer state at epoch {} to {}".format(
+    #         iteration, checkpoint_path
+    #     )
+    # )
     if hasattr(model, "module"):
         state_dict = model.module.state_dict()
     else:
         state_dict = model.state_dict()
+    if delete_old:
+        latest_checkpoint = latest_checkpoint_path(checkpoint_path, regex=("G_*.pth" if checkpoint_type.startswith("G") else "D_*.pth"))
+
     torch.save(
         {
             "model": state_dict,
             "iteration": iteration,
             "optimizer": optimizer.state_dict(),
             "learning_rate": learning_rate,
         },
-        checkpoint_path,
+        os.path.join(checkpoint_path, checkpoint_type),
     )
+    # delete after saving new checkpoint to avoid loss if save fails
+    if delete_old and latest_checkpoint is not None:
+        os.remove(latest_checkpoint)
 
 
 def save_checkpoint_d(combd, sbd, optimizer, learning_rate, iteration, checkpoint_path):
-    logger.info(
-        "Saving model and optimizer state at epoch {} to {}".format(
-            iteration, checkpoint_path
-        )
-    )
+    # logger.info(
+    #     "Saving model and optimizer state at epoch {} to {}".format(
+    #         iteration, checkpoint_path
+    #     )
+
     if hasattr(combd, "module"):
         state_dict_combd = combd.module.state_dict()
     else:
@@ -204,7 +200,7 @@ def latest_checkpoint_path(dir_path, regex="G_*.pth"):
     f_list = glob.glob(os.path.join(dir_path, regex))
     f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
     x = f_list[-1]
-    print(x)
+    # print(x)
     return x
 
 
@@ -247,9 +243,7 @@ def plot_alignment_to_numpy(alignment, info=None):
     import numpy as np
 
     fig, ax = plt.subplots(figsize=(6, 4))
-    im = ax.imshow(
-        alignment.transpose(), aspect="auto", origin="lower", interpolation="none"
-    )
+    im = ax.imshow(alignment.transpose(), aspect="auto", origin="lower", interpolation="none")
     fig.colorbar(im, ax=ax)
     xlabel = "Decoder timestep"
     if info is not None:
@@ -302,35 +296,21 @@ def get_hparams(init=True):
         required=True,
         help="checkpoint save frequency (epoch)",
     )
-    parser.add_argument(
-        "-te", "--total_epoch", type=int, required=True, help="total_epoch"
-    )
-    parser.add_argument(
-        "-pg", "--pretrainG", type=str, default="", help="Pretrained Discriminator path"
-    )
-    parser.add_argument(
-        "-pd", "--pretrainD", type=str, default="", help="Pretrained Generator path"
-    )
+    parser.add_argument("-te", "--total_epoch", type=int, required=True, help="total_epoch")
+    parser.add_argument("-pg", "--pretrainG", type=str, default="", help="Pretrained Discriminator path")
+    parser.add_argument("-pd", "--pretrainD", type=str, default="", help="Pretrained Generator path")
     parser.add_argument("-g", "--gpus", type=str, default="0", help="split by -")
-    parser.add_argument(
-        "-bs", "--batch_size", type=int, required=True, help="batch size"
-    )
-    parser.add_argument(
-        "-e", "--experiment_dir", type=str, required=True, help="experiment dir"
-    )  # -m
-    parser.add_argument(
-        "-sr", "--sample_rate", type=str, required=True, help="sample rate, 32k/40k/48k"
-    )
+    parser.add_argument("-bs", "--batch_size", type=int, required=True, help="batch size")
+    parser.add_argument("-e", "--experiment_dir", type=str, required=True, help="experiment dir")  # -m
+    parser.add_argument("-sr", "--sample_rate", type=str, required=True, help="sample rate, 32k/40k/48k")
     parser.add_argument(
         "-sw",
         "--save_every_weights",
         type=str,
         default="0",
         help="save the extracted model in weights directory when saving checkpoints",
     )
-    parser.add_argument(
-        "-v", "--version", type=str, required=True, help="model version"
-    )
+    parser.add_argument("-v", "--version", type=str, required=True, help="model version")
     parser.add_argument(
         "-f0",
         "--if_f0",
@@ -414,11 +394,7 @@ def get_hparams_from_file(config_path):
 def check_git_hash(model_dir):
     source_dir = os.path.dirname(os.path.realpath(__file__))
     if not os.path.exists(os.path.join(source_dir, ".git")):
-        logger.warn(
-            "{} is not a git repository, therefore hash value comparison will be ignored.".format(
-                source_dir
-            )
-        )
+        logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(source_dir))
         return
 
     cur_hash = subprocess.getoutput("git rev-parse HEAD")
@@ -427,11 +403,7 @@ def check_git_hash(model_dir):
     if os.path.exists(path):
         saved_hash = open(path).read()
         if saved_hash != cur_hash:
-            logger.warn(
-                "git hash values are different. {}(saved) != {}(current)".format(
-                    saved_hash[:8], cur_hash[:8]
-                )
-            )
+            logger.warn("git hash values are different. {}(saved) != {}(current)".format(saved_hash[:8], cur_hash[:8]))
     else:
         open(path, "w").write(cur_hash)
 

train_nsf_sim_cache_sid_load_pretrain.py

@@ -98,9 +98,7 @@ def run(rank, n_gpus, hps):
         writer = SummaryWriter(log_dir=hps.model_dir)
         writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
 
-    dist.init_process_group(
-        backend="gloo", init_method="env://", world_size=n_gpus, rank=rank
-    )
+    dist.init_process_group(backend="gloo", init_method="env://", world_size=n_gpus, rank=rank)
     torch.manual_seed(hps.train.seed)
     if torch.cuda.is_available():
         torch.cuda.set_device(rank)
@@ -176,93 +174,55 @@ def run(rank, n_gpus, hps):
         net_d = DDP(net_d)
 
     try:  # 如果能加载自动resume
-        _, _, _, epoch_str = utils.load_checkpoint(
-            utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_d
-        )  # D多半加载没事
+        _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_d)  # D多半加载没事
         if rank == 0:
             logger.info("loaded D")
         # _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g,load_opt=0)
-        _, _, _, epoch_str = utils.load_checkpoint(
-            utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g
-        )
+        _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g)
+
+        scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)
+        scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)
 
-        scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
-            optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2
-        )
-        scheduler_d = torch.optim.lr_scheduler.ExponentialLR(
-            optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2
-        )
-        
         # epoch_str = 1
         # global_step = 0
     except:  # 如果首次不能加载，加载pretrain
         # traceback.print_exc()
         epoch_str = 0
         global_step = 0
         if hps.pretrainG != "":
-
             if rank == 0:
                 logger.info("loaded pretrained %s" % (hps.pretrainG))
-            print(
-                net_g.module.load_state_dict(
-                    torch.load(hps.pretrainG, map_location="cpu")["model"]
-                )
-            )  ##测试不加载优化器
+            print(net_g.module.load_state_dict(torch.load(hps.pretrainG, map_location="cpu")["model"]))  ##测试不加载优化器
         if hps.pretrainD != "":
-
             if rank == 0:
                 logger.info("loaded pretrained %s" % (hps.pretrainD))
-            print(
-                net_d.module.load_state_dict(
-                    torch.load(hps.pretrainD, map_location="cpu")["model"]
-                )
-            )
-            scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
-                optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 1
-            )
-            scheduler_d = torch.optim.lr_scheduler.ExponentialLR(
-                optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 1
-            )
+            print(net_d.module.load_state_dict(torch.load(hps.pretrainD, map_location="cpu")["model"]))
+            scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 1)
+            scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 1)
     global_step = epoch_str * len(train_loader)
 
-    
-
     scaler = GradScaler(enabled=hps.train.fp16_run)
 
     cache = []
     # for epoch in tqdm.tqdm(range(epoch_str, hps.train.epochs + 1), desc="Training progress", position=1, leave=True):
-    with tqdm.tqdm(total=(hps.total_epoch), desc=f"Training progress, last ckpt saved at epoch: {epoch_str}", position=0, leave=True, initial=0 if global_step == 0 else epoch_str, dynamic_ncols=True) as pbar:
+
+    # disable if rank is not 0 to avoid duplicate main progress bars on each process
+    with tqdm.tqdm(
+        total=(hps.total_epoch),
+        desc=f"Training progress, last ckpt saved at epoch: {epoch_str}",
+        position=0,
+        leave=True,
+        initial=0 if global_step == 0 else epoch_str,
+        dynamic_ncols=True,
+        disable=(rank != 0),
+    ) as pbar:
         for epoch in range(epoch_str + 1, hps.train.epochs + 1):
             if rank == 0:
                 train_and_evaluate(
-                    rank,
-                    epoch,
-                    hps,
-                    [net_g, net_d],
-                    [optim_g, optim_d],
-                    [scheduler_g, scheduler_d],
-                    scaler,
-                    [train_loader, None],
-                    logger,
-                    [writer, writer_eval],
-                    cache,
-                    pbar
+                    rank, epoch, hps, [net_g, net_d], [optim_g, optim_d], [scheduler_g, scheduler_d], scaler, [train_loader, None], logger, [writer, writer_eval], cache, pbar
                 )
             else:
-                train_and_evaluate(
-                    rank,
-                    epoch,
-                    hps,
-                    [net_g, net_d],
-                    [optim_g, optim_d],
-                    [scheduler_g, scheduler_d],
-                    scaler,
-                    [train_loader, None],
-                    None,
-                    None,
-                    cache,
-                    pbar
-                )
+                train_and_evaluate(rank, epoch, hps, [net_g, net_d], [optim_g, optim_d], [scheduler_g, scheduler_d], scaler, [train_loader, None], None, None, cache, pbar)
             pbar.update(1)
             scheduler_g.step()
             scheduler_d.step()
@@ -540,36 +500,25 @@ def train_and_evaluate(
         # /Run steps
 
     if epoch % hps.save_every_epoch == 0 and rank == 0:
-        if hps.if_latest == 0:
-            utils.save_checkpoint(
-                net_g,
-                optim_g,
-                hps.train.learning_rate,
-                epoch,
-                os.path.join(hps.model_dir, "G_{}.pth".format(global_step)),
-            )
-            utils.save_checkpoint(
-                net_d,
-                optim_d,
-                hps.train.learning_rate,
-                epoch,
-                os.path.join(hps.model_dir, "D_{}.pth".format(global_step)),
-            )
-        else:
-            utils.save_checkpoint(
-                net_g,
-                optim_g,
-                hps.train.learning_rate,
-                epoch,
-                os.path.join(hps.model_dir, "G_{}.pth".format(2333333)),
-            )
-            utils.save_checkpoint(
-                net_d,
-                optim_d,
-                hps.train.learning_rate,
-                epoch,
-                os.path.join(hps.model_dir, "D_{}.pth".format(2333333)),
-            )
+        utils.save_checkpoint(
+            net_g,
+            optim_g,
+            hps.train.learning_rate,
+            epoch,
+            hps.model_dir,
+            "G_{}.pth".format(global_step),
+            True if hps.if_latest == 1 else False,
+        )
+        utils.save_checkpoint(
+            net_d,
+            optim_d,
+            hps.train.learning_rate,
+            epoch,
+            hps.model_dir,
+            "D_{}.pth".format(global_step),
+            True if hps.if_latest == 1 else False,
+        )
+
         pbar.set_description(f"Training progress, last ckpt saved at epoch {epoch}")
         if rank == 0 and hps.save_every_weights == "1":
             if hasattr(net_g, "module"):
@@ -603,14 +552,7 @@ def train_and_evaluate(
             ckpt = net_g.module.state_dict()
         else:
             ckpt = net_g.state_dict()
-        logger.info(
-            "saving final ckpt:%s"
-            % (
-                savee(
-                    ckpt, hps.sample_rate, hps.if_f0, hps.name, epoch, hps.version, hps
-                )
-            )
-        )
+        logger.info("saving final ckpt:%s" % (savee(ckpt, hps.sample_rate, hps.if_f0, hps.name, epoch, hps.version, hps)))
         sleep(1)
         os._exit(2333333)