Set 2n rule..

This commit is contained in:
Daniel Povey 2022-04-09 18:15:56 +08:00
parent 0f8ee68af2
commit db72aee1f0

View File

@ -154,15 +154,17 @@ def get_parser():
parser.add_argument( parser.add_argument(
"--lr-begin-steps", "--lr-begin-steps",
type=float, type=float,
default=20000, default=25000,
help="Number of steps that affects how rapidly the learning rate initially decreases" help="Number of steps that affects how rapidly the learning rate initially decreases"
) )
parser.add_argument( parser.add_argument(
"--lr-end-epochs", "--lr-end-epochs",
type=float, type=float,
default=10, default=-1,
help="Number of epochs that affects how rapidly the learning rate finally decreases" help="""Number of epochs that affects how rapidly the learning rate finally decreases;
if -1, will be set the same as --num-epochs
"""
) )
parser.add_argument( parser.add_argument(
@ -783,12 +785,14 @@ def run(rank, world_size, args):
lr=params.initial_lr) lr=params.initial_lr)
# The `epoch` variable in the lambda expression binds to the value below # The `epoch` variable in the lambda expression binds to the value below
# in `for epoch in range(params.start_epoch, params.num_epochs):`. # in `for epoch in range(params.start_epoch, params.num_epochs):`. But set it to 0
# here to avoid crash in constructor.
epoch = 0 epoch = 0
lr_end_epochs = params.lr_end_epochs if params.lr_end_epochs > 0 else params.num_epochs
scheduler = torch.optim.lr_scheduler.LambdaLR( scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer, optimizer,
lambda step: (((step + params.lr_begin_steps) / params.lr_begin_steps) ** -0.5 * lambda step: (((step + params.lr_begin_steps) / params.lr_begin_steps) ** -0.5 *
math.exp(-epoch / params.lr_end_epochs))) ((epoch + lr_end_epochs) / lr_end_epochs) ** -2.0))
if checkpoints and "optimizer" in checkpoints: if checkpoints and "optimizer" in checkpoints: