From 576424fe583a5710329d6feed573fef2f1446e43 Mon Sep 17 00:00:00 2001 From: Emberstar <969242373@qq.com> Date: Wed, 14 Jun 2023 10:18:19 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E5=8A=A0=E8=BD=BDlatest=20mo?= =?UTF-8?q?del=E7=9A=84=E6=96=B9=E5=BC=8F=EF=BC=8C=E4=BF=AE=E6=94=B9global?= =?UTF-8?q?=5Fstep=E8=AE=A1=E7=AE=97=EF=BC=8C=E5=A2=9E=E5=8A=A0preserved?= =?UTF-8?q?=E5=8F=82=E6=95=B0=EF=BC=8C=E5=A2=9E=E5=8A=A0train=5Fwith=5Fpre?= =?UTF-8?q?trained=5Fmodel=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- finetune_speaker_v2.py | 33 ++++++++++++++++++++++----------- utils.py | 27 ++++++++++++++++++++++++--- 2 files changed, 46 insertions(+), 14 deletions(-) diff --git a/finetune_speaker_v2.py b/finetune_speaker_v2.py index a0709d4..75baf7e 100644 --- a/finetune_speaker_v2.py +++ b/finetune_speaker_v2.py @@ -100,18 +100,26 @@ def run(rank, n_gpus, hps): # load existing model if hps.cont: try: - _, _, _, epoch_str = utils.load_checkpoint("./OUTPUT_MODEL/G_latest.pth", net_g, None) - _, _, _, epoch_str = utils.load_checkpoint("./OUTPUT_MODEL/D_latest.pth", net_d, None) - global_step = epoch_str * hps.train.batch_size + _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_[0-9]*.pth"), net_g, None) + _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_[0-9]*.pth"), net_d, None) + global_step = (epoch_str - 1) * len(train_loader) except: print("Failed to find latest checkpoint, loading G_0.pth...") - _, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/G_0.pth", net_g, None) - _, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/D_0.pth", net_d, None) + if hps.train_with_pretrained_model: + print("Train with pretrained model...") + _, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/G_0.pth", net_g, None) + _, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/D_0.pth", net_d, None) + else: + print("Train without pretrained model...") epoch_str = 1 global_step = 0 else: - _, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/G_0.pth", net_g, None) - _, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/D_0.pth", net_d, None) + if hps.train_with_pretrained_model: + print("Train with pretrained model...") + _, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/G_0.pth", net_g, None) + _, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/D_0.pth", net_d, None) + else: + print("Train without pretrained model...") epoch_str = 1 global_step = 0 # freeze all other layers except speaker embedding @@ -256,13 +264,16 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade utils.save_checkpoint(net_g, None, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "G_latest.pth".format(global_step))) utils.save_checkpoint(net_d, None, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "D_{}.pth".format(global_step))) - utils.save_checkpoint(net_d, None, hps.train.learning_rate, epoch, - os.path.join(hps.model_dir, "D_latest.pth".format(global_step))) - old_g=os.path.join(hps.model_dir, "G_{}.pth".format(global_step-4000)) - old_d=os.path.join(hps.model_dir, "D_{}.pth".format(global_step-4000)) + # utils.save_checkpoint(net_d, None, hps.train.learning_rate, epoch, + # os.path.join(hps.model_dir, "D_latest.pth".format(global_step))) + old_g = utils.oldest_checkpoint_path(hps.model_dir, "G_[0-9]*.pth", + preserved=hps.preserved) # Preserve 4 (default) historical checkpoints. + old_d = utils.oldest_checkpoint_path(hps.model_dir, "D_[0-9]*.pth", preserved=hps.preserved) if os.path.exists(old_g): + print(f"remove {old_g}") os.remove(old_g) if os.path.exists(old_d): + print(f"remove {old_d}") os.remove(old_d) global_step += 1 if epoch > hps.max_epochs: diff --git a/utils.py b/utils.py index 51e3dee..205d929 100644 --- a/utils.py +++ b/utils.py @@ -204,14 +204,29 @@ def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios= writer.add_audio(k, v, global_step, audio_sampling_rate) -def latest_checkpoint_path(dir_path, regex="G_*.pth"): +def extract_digits(f): + digits = "".join(filter(str.isdigit, f)) + return int(digits) if digits else -1 + + +def latest_checkpoint_path(dir_path, regex="G_[0-9]*.pth"): f_list = glob.glob(os.path.join(dir_path, regex)) - f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f)))) + f_list.sort(key=lambda f: extract_digits(f)) x = f_list[-1] - print(x) + print(f"latest_checkpoint_path:{x}") return x +def oldest_checkpoint_path(dir_path, regex="G_[0-9]*.pth", preserved=4): + f_list = glob.glob(os.path.join(dir_path, regex)) + f_list.sort(key=lambda f: extract_digits(f)) + if len(f_list) > preserved: + x = f_list[0] + print(f"oldest_checkpoint_path:{x}") + return x + return "" + + def plot_spectrogram_to_numpy(spectrogram): global MATPLOTLIB_FLAG if not MATPLOTLIB_FLAG: @@ -288,6 +303,10 @@ def get_hparams(init=True): help='finetune epochs') parser.add_argument('--cont', type=bool, default=False, help='whether to continue training on the latest checkpoint') parser.add_argument('--drop_speaker_embed', type=bool, default=False, help='whether to drop existing characters') + parser.add_argument('--train_with_pretrained_model', type=bool, default=True, + help='whether to train with pretrained model') + parser.add_argument('--preserved', type=int, default=4, + help='Number of preserved models') args = parser.parse_args() model_dir = os.path.join("./", args.model) @@ -312,6 +331,8 @@ def get_hparams(init=True): hparams.max_epochs = args.max_epochs hparams.cont = args.cont hparams.drop_speaker_embed = args.drop_speaker_embed + hparams.train_with_pretrained_model = args.train_with_pretrained_model + hparams.preserved = args.preserved return hparams