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