def _convert_to_v2_lr_schedules(params):
decay_type = params.get("decay_type")
if not decay_type:
return
try:
# Check if the learning rate schedule exists in V2.
schedules_lib.get_lr_schedule_class(decay_type)
except ValueError:
if decay_type.startswith("noam_decay"):
params["decay_type"] = "NoamDecay"
if "decay_params" not in params:
model_dim = _delete_opt(params, "decay_rate")
warmup_steps = _delete_opt(params, "decay_steps")
params["decay_params"] = dict(model_dim=model_dim,
warmup_steps=warmup_steps)
else:
raise ValueError(
"params/decay_type should be manually converted: no registered "
"conversion for decay type %s" % decay_type)
def testGetScheduleClass(self):
with self.assertRaises(ValueError):
lr_schedules.get_lr_schedule_class("ScheduleWrapper")
lr_schedules.get_lr_schedule_class(
"NoamDecay") == lr_schedules.NoamDecay