Added capability of continue training from previous checkpoints

This commit is contained in:
Plachta
2023-06-12 19:20:32 +08:00
parent 3b8d7b5ef4
commit 9bbc9e9246
+3 -4
View File
@@ -100,10 +100,9 @@ def run(rank, n_gpus, hps):
# load existing model
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
_, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/G_latest.pth", net_g, None)
_, _, _, epoch_str = utils.load_checkpoint("./pretrained_models/D_latest.pth", net_d, None)
global_step = epoch_str * hps.train.batch_size
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)