diff --git a/LOCAL.md b/LOCAL.md index 931b1bf..3008ebe 100644 --- a/LOCAL.md +++ b/LOCAL.md @@ -96,6 +96,7 @@ 11. Start Training. Run `python finetune_speaker_v2.py -m "./OUTPUT_MODEL" --max_epochs "{Maximum_epochs}" --drop_speaker_embed True` Do replace `{Maximum_epochs}` with your desired number of epochs. Empirically, 100 or more is recommended. + To continue training on previous checkpoint, change the training command to: `python finetune_speaker_v2.py -m "./OUTPUT_MODEL" --max_epochs "{Maximum_epochs}" --drop_speaker_embed True --cont True`. Before you do this, make sure you have previous `G_latest.pth` and `D_latest.pth` under `./OUTPUT_MODEL/` directory. To view training progress, open a new terminal and `cd` to the project root directory, run `tensorboard --logdir="./OUTPUT_MODEL"`, then visit `localhost:6006` with your web browser. 12. After training is completed, you can use your model by running: `python VC_inference.py --model_dir ./OUTPUT_MODEL/G_latest.pth --share True` diff --git a/finetune_speaker_v2.py b/finetune_speaker_v2.py index 85fa044..809000e 100644 --- a/finetune_speaker_v2.py +++ b/finetune_speaker_v2.py @@ -98,8 +98,17 @@ def run(rank, n_gpus, hps): net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(rank) # load existing model - _, _, _, _ = utils.load_checkpoint("./pretrained_models/G_0.pth", net_g, None, drop_speaker_emb=hps.drop_speaker_embed) - _, _, _, _ = utils.load_checkpoint("./pretrained_models/D_0.pth", net_d, None) + G_ckpt = "./pretrained_models/G_latest.pth" if hps.cont else "./pretrained_models/G_0.pth" + D_ckpt = "./pretrained_models/D_latest.pth" if hps.cont else "./pretrained_models/D_0.pth" + try: + _, _, _, _ = utils.load_checkpoint(G_ckpt, net_g, None, + drop_speaker_emb=hps.drop_speaker_embed) + except Exception: + _, _, _, _ = utils.load_checkpoint("./pretrained_models/G_0.pth", net_g, None, drop_speaker_emb=hps.drop_speaker_embed) + try: + _, _, _, _ = utils.load_checkpoint(D_ckpt, net_d, None) + except Exception: + _, _, _, _ = utils.load_checkpoint("./pretrained_models/D_0.pth", net_d, None) epoch_str = 1 global_step = 0 # freeze all other layers except speaker embedding @@ -243,13 +252,15 @@ 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_{}.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, optim_d, 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_{}.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-400)) + 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 os.path.exists(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 d72ea0f..954db6a 100644 --- a/utils.py +++ b/utils.py @@ -286,6 +286,7 @@ 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') args = parser.parse_args()