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