From 3b8d7b5ef4f61b6d2b440205ae834efb3fe68cf7 Mon Sep 17 00:00:00 2001
From: Plachta <liusongting07@gmail.com>
Date: Mon, 12 Jun 2023 19:11:52 +0800
Subject: [PATCH] Added capability of continue training from previous
 checkpoints

---
 finetune_speaker_v2.py | 29 ++++++++++++++++-------------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/finetune_speaker_v2.py b/finetune_speaker_v2.py
index 809000e..5afbae1 100644
--- a/finetune_speaker_v2.py
+++ b/finetune_speaker_v2.py
@@ -98,19 +98,22 @@ def run(rank, n_gpus, hps):
   net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(rank)
 
   # load existing model
-  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
+  if hps.cont:
+      try:
+          _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, None)
+          _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, None)
+          epoch_str = 1
+          global_step = 0
+      except:
+          _, _, _, 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)
+          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)
+      epoch_str = 1
+      global_step = 0
   # freeze all other layers except speaker embedding
   for p in net_g.parameters():
       p.requires_grad = True