tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Fix error on resume training: KeyError: "param 'initial_lr' is not specified in param_groups[0] when resuming an optimizer"

Browse Source
This commit is contained in:
Steve Wang
2023-09-08 19:17:39 +08:00
parent f7db8c997a
commit 1791d9240a
1 changed files with 16 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
train_ms.py
View File

@@ -157,12 +157,20 @@ def run(rank, n_gpus, hps):
net_dur_disc = DDP(net_dur_disc, device_ids=[rank], find_unused_parameters=True) net_dur_disc = DDP(net_dur_disc, device_ids=[rank], find_unused_parameters=True)
try: try:
if net_dur_disc is not None: if net_dur_disc is not None:
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "DUR_*.pth"), net_dur_disc, optim_dur_disc, skip_optimizer=True) _, _, dur_resume_lr, epoch_str = utils.load_checkpoint(
_, optim_g, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, utils.latest_checkpoint_path(hps.model_dir, "DUR_*.pth"), net_dur_disc, optim_dur_disc,
optim_g, skip_optimizer=True) skip_optimizer=True)
_, optim_d, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, _, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint(
optim_d, skip_optimizer=True) utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g,
optim_g, skip_optimizer=True)
_, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint(
utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d,
optim_d, skip_optimizer=True)
if not optim_g.param_groups[0].get("initial_lr"):
optim_g.param_groups[0]["initial_lr"] = g_resume_lr
if not optim_d.param_groups[0].get("initial_lr"):
optim_d.param_groups[0]["initial_lr"] = d_resume_lr
epoch_str = max(epoch_str, 1) epoch_str = max(epoch_str, 1)
global_step = (epoch_str - 1) * len(train_loader) global_step = (epoch_str - 1) * len(train_loader)
except Exception as e: except Exception as e:
@@ -174,6 +182,8 @@ def run(rank, n_gpus, hps):
scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2) scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)
scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2) scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)
if net_dur_disc is not None: if net_dur_disc is not None:
if not optim_dur_disc.param_groups[0].get("initial_lr"):
optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
scheduler_dur_disc = torch.optim.lr_scheduler.ExponentialLR(optim_dur_disc, gamma=hps.train.lr_decay, last_epoch=epoch_str-2) scheduler_dur_disc = torch.optim.lr_scheduler.ExponentialLR(optim_dur_disc, gamma=hps.train.lr_decay, last_epoch=epoch_str-2)
else: else:
scheduler_dur_disc = None scheduler_dur_disc = None