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

Merge pull request #21 from kale4eat/dev-warmup

Browse Source 
エポックごとのスケジューリングによるウォームアップ
This commit is contained in:
litagin02
2024-01-09 15:24:26 +09:00
committed by GitHub
parent 59f5b9446e 7d854431f3
commit d461b2e78d
1 changed files with 14 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
train_ms.py
View File

@@ -356,19 +356,18 @@ def run():
epoch_str = 1 epoch_str = 1
global_step = 0 global_step = 0
steps_per_epoch = len(train_loader) def lr_lambda(epoch):
warmup_steps = hps.train.warmup_epochs * steps_per_epoch """
Learning rate scheduler for warmup and exponential decay.
def lr_lambda(step): - During the warmup period, the learning rate increases linearly.
"""for warmup""" - After the warmup period, the learning rate decreases exponentially.
if step < warmup_steps: """
return step / warmup_steps if epoch < hps.train.warmup_epochs:
return float(epoch) / float(max(1, hps.train.warmup_epochs))
else: else:
return hps.train.lr_decay ** (step - warmup_steps) return hps.train.lr_decay ** (epoch - hps.train.warmup_epochs)
scheduler_last_epoch = -1 scheduler_last_epoch = epoch_str - 2
if global_step > 0:
scheduler_last_epoch = global_step - 1
scheduler_g = torch.optim.lr_scheduler.LambdaLR( scheduler_g = torch.optim.lr_scheduler.LambdaLR(
optim_g, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch optim_g, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch
) )
@@ -376,12 +375,6 @@ def run():
optim_d, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch optim_d, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch
) )
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"):
if global_step == 0:
optim_dur_disc.param_groups[0]["initial_lr"] = hps.train.learning_rate
else:
optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
initial_lr = optim_dur_disc.param_groups[0].get("initial_lr")
scheduler_dur_disc = torch.optim.lr_scheduler.LambdaLR( scheduler_dur_disc = torch.optim.lr_scheduler.LambdaLR(
optim_dur_disc, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch optim_dur_disc, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch
) )
@@ -436,6 +429,10 @@ def run():
pbar, pbar,
initial_step, initial_step,
) )
scheduler_g.step()
scheduler_d.step()
if net_dur_disc is not None:
scheduler_dur_disc.step()
if epoch == hps.train.epochs: if epoch == hps.train.epochs:
# Save the final models # Save the final models
@@ -752,10 +749,6 @@ def train_and_evaluate(
for_infer=True, for_infer=True,
) )
scheduler_g.step()
scheduler_d.step()
if net_dur_disc is not None:
scheduler_dur_disc.step()
global_step += 1 global_step += 1
if pbar is not None: if pbar is not None:
pbar.set_description( pbar.set_description(