diff --git a/finetune_speaker_v2.py b/finetune_speaker_v2.py index a0709d4..5fc4fd8 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_latest.pth"), net_g, None) + _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_latest.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 @@ -252,18 +260,26 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade if global_step % hps.train.eval_interval == 0: evaluate(hps, net_g, eval_loader, writer_eval) - utils.save_checkpoint(net_g, None, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "G_{}.pth".format(global_step))) + 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)) - if os.path.exists(old_g): - os.remove(old_g) - if os.path.exists(old_d): - os.remove(old_d) + if hps.preserved > 0: + utils.save_checkpoint(net_g, None, hps.train.learning_rate, epoch, + os.path.join(hps.model_dir, "G_{}.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))) + 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: print("Maximum epoch reached, closing training...") diff --git a/utils.py b/utils.py index 51e3dee..648ae27 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: @@ -278,6 +293,17 @@ def load_filepaths_and_text(filename, split="|"): return filepaths_and_text +def str2bool(v): + if isinstance(v, bool): + return v + if v.lower() in ('yes', 'true', 't', 'y', '1'): + return True + elif v.lower() in ('no', 'false', 'f', 'n', '0'): + return False + else: + raise argparse.ArgumentTypeError('Boolean value expected.') + + def get_hparams(init=True): parser = argparse.ArgumentParser() parser.add_argument('-c', '--config', type=str, default="./configs/modified_finetune_speaker.json", @@ -286,8 +312,12 @@ def get_hparams(init=True): help='Model name') parser.add_argument('-n', '--max_epochs', type=int, default=50, 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('--cont', type=str2bool, default=False, help='whether to continue training on the latest checkpoint') + parser.add_argument('--drop_speaker_embed', type=str2bool, default=False, help='whether to drop existing characters') + parser.add_argument('--train_with_pretrained_model', type=str2bool, 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 +342,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