Change exponential part of lrate to be epoch based

This commit is contained in:
Daniel Povey 2022-04-08 16:24:21 +08:00
parent 6ee32cf7af
commit f587cd527d

View File

@ -27,10 +27,8 @@ export CUDA_VISIBLE_DEVICES="0,1,2,3"
--start-epoch 0 \ --start-epoch 0 \
--exp-dir pruned_transducer_stateless2/exp \ --exp-dir pruned_transducer_stateless2/exp \
--full-libri 1 \ --full-libri 1 \
--max-duration 300 \ --max-duration 300
--initial-lr 0.003 \
--lr-begin-steps 20000 \
--lr-end-steps 50000
""" """
@ -161,10 +159,10 @@ def get_parser():
) )
parser.add_argument( parser.add_argument(
"--lr-end-steps", "--lr-end-epochs",
type=float, type=float,
default=50000, default=10,
help="Number of steps that affects how rapidly the learning rate finally decreases" help="Number of epochs that affects how rapidly the learning rate finally decreases"
) )
parser.add_argument( parser.add_argument(
@ -783,10 +781,13 @@ def run(rank, world_size, args):
optimizer = Eve( optimizer = Eve(
model.parameters(), model.parameters(),
lr=params.initial_lr) lr=params.initial_lr)
# The `epoch` variable in the lambda expression binds to the value below
# in `for epoch in range(params.start_epoch, 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(-step / params.lr_end_steps))) math.exp(-epoch / params.lr_end_epochs)))
if checkpoints and "optimizer" in checkpoints: if checkpoints and "optimizer" in checkpoints: