提交 820b4db2 编写于 作者: H Hui Zhang

with all args for scheduler

上级 c4da9a7f
......@@ -315,18 +315,11 @@ class U2Trainer(Trainer):
scheduler_conf = train_config.scheduler_conf
scheduler_args = {
"learning_rate":
optim_conf.lr,
"verbose":
False,
"warmup_steps":
scheduler_conf.warmup_steps
if "warmup_steps" in scheduler_conf else None,
"gamma":
scheduler_conf.lr_decay if "lr_decay" in scheduler_conf else None,
"d_model":
model_conf.encoder_conf.output_size
if scheduler_type == "noam" else None,
"learning_rate": optim_conf.lr,
"verbose": False,
"warmup_steps": scheduler_conf.warmup_steps,
"gamma": scheduler_conf.lr_decay,
"d_model": model_conf.encoder_conf.output_size,
}
lr_scheduler = LRSchedulerFactory.from_args(scheduler_type,
scheduler_args)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册