Merge pull request #141 from ReDeiPirati/yellowfin_optimizer

Integrate YellowFin Optimizer(with test) in T2T

Author: lukaszkaiser

tensor2tensor/models/common_hparams.py

@@ -202,4 +202,9 @@ def basic_range1(ranged_hparams):
   rhp.set_float("optimizer_adam_beta1", 0.8, 0.9)
   rhp.set_float("optimizer_adam_beta2", 0.995, 0.999)
   rhp.set_categorical("optimizer",
-                      ["Adam", "Adagrad", "Momentum", "RMSProp", "SGD"])
+                      ["Adam",
+                       "Adagrad",
+                       "Momentum",
+                       "RMSProp",
+                       "SGD",
+                       "YellowFin"])

tensor2tensor/utils/trainer_utils.py

@@ -43,6 +43,7 @@
 import tensorflow as tf
 from tensorflow.contrib.learn.python.learn import learn_runner
 from tensorflow.python.ops import init_ops
+from tensor2tensor.utils.yellowfin import YellowFinOptimizer
 
 # Number of samples to draw for an image input (in such cases as captioning)
 IMAGE_DECODE_LENGTH = 100
@@ -1141,6 +1142,10 @@ def __init__(self, optimizer_name, lr, hparams):
     elif optimizer_name == "Momentum":
       self._opt = tf.train.MomentumOptimizer(
           lr, momentum=hparams.optimizer_momentum_momentum)
+    elif optimizer_name == "YellowFin":
+      tf.logging.info("Init YellowFin Optimizer.")
+      self._opt = YellowFinOptimizer(
+          learning_rate=lr, momentum=hparams.optimizer_momentum_momentum)
     else:
       self._opt = tf.contrib.layers.OPTIMIZER_CLS_NAMES[optimizer_name](lr)