Revert "noam" learning-rate-scheme to use linear warmup. Add learning_rate_schedule hparam to specify a schedule that does not have separate warmup and decay phases.

PiperOrigin-RevId: 185042750

Author: nshazeer

tensor2tensor/layers/common_hparams.py

@@ -63,6 +63,11 @@ def basic_params1():
       optimizer_momentum_nesterov=False,
       weight_decay=1e-6,
       weight_noise=0.0,
+      learning_rate_schedule="warmup_and_decay",
+      # If learning_rate_schedule=="warmup_and_decay", then this specifies
+      # the decay part of the schedule.
+      # The warmup is always exponential.
+      # TODO(noam): add a hyperparameter to control the warmup.
       learning_rate_decay_scheme="none",
       # decay_steps and decay_staircase for learning_rate_decay_scheme=="exp"
       learning_rate_decay_steps=5000,

tensor2tensor/models/transformer.py

@@ -877,7 +877,7 @@ def transformer_base_v1():
   hparams.max_length = 256
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
   hparams.optimizer_adam_epsilon = 1e-9
-  hparams.learning_rate_decay_scheme = "noam"
+  hparams.learning_rate_schedule = "linear_warmup_rsqrt_decay"
   hparams.learning_rate = 0.1
   hparams.learning_rate_warmup_steps = 4000
   hparams.initializer_gain = 1.0

tensor2tensor/utils/optimize.py

@@ -168,10 +168,6 @@ def learning_rate_decay(hparams, warmup_steps=0):
                                    hparams.learning_rate_boundaries,
                                    hparams.learning_rate_multiples)
 
-  if scheme == "noam":
-    return 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
-        (global_step + 1) * warmup_steps**-1.5, (global_step + 1)**-0.5)
-
   if scheme == "cosine":
     cycle_steps = hparams.learning_rate_cosine_cycle_steps
     cycle_position = global_step % (2 * cycle_steps)
@@ -224,6 +220,23 @@ def learning_rate_decay_with_warmup(hparams, num_worker_replicas=1):
   return tf.where(global_step < warmup_steps, warmup, decay)
 
 
+def learning_rate_schedule(hparams, num_worker_replicas=1):
+  """Learning rate schedule based on hparams."""
+  schedule = hparams.learning_rate_schedule
+  warmup_steps = tf.to_float(hparams.learning_rate_warmup_steps)
+  global_step = tf.to_float(tf.train.get_or_create_global_step())
+  if hparams.learning_rate_decay_scheme == "noam":
+    # backwards compatiblity with previous behavior
+    schedule = "linear_warmup_rsqrt_decay"
+  if schedule == "warmup_and_decay":
+    return learning_rate_decay_with_warmup(hparams, num_worker_replicas)
+  elif schedule == "linear_warmup_rsqrt_decay":
+    return 5000.0 * hparams.hidden_size**-0.5 * tf.minimum(
+        (global_step + 1) * warmup_steps**-1.5, (global_step + 1)**-0.5)
+  else:
+    raise ValueError("Unrecognized learning rate schedule: %s" % schedule)
+
+
 def weight_decay_and_noise(loss, hparams, learning_rate, var_list=None):
   """Apply weight decay and weight noise."""
   if var_list is None:

tensor2tensor/utils/t2t_model.py

@@ -296,7 +296,7 @@ def optimize(self, loss, num_async_replicas=1):
     """Return a training op minimizing loss."""
     tf.logging.info("Base learning rate: %f", self.hparams.learning_rate)
     lr = self.hparams.learning_rate
-    decay_rate = optimize.learning_rate_decay_with_warmup(self.hparams)
+    decay_rate = optimize.learning_rate_schedule(self.hparams)
     lr *= decay_rate
     if self.hparams.learning_rate_minimum:
       lr_min = float(self.hparams.learning_rate_minimum)