grassmann_optimizer.py

#from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow.python.training import optimizer

from gutils import unit
from gutils import gproj
from gutils import clip_by_norm
from gutils import xTy
from gutils import gexp
from gutils import gpt
from gutils import gpt2

class HybridOptimizer(optimizer.Optimizer):
  def __init__(self, optimizer_a, optimizer_b, use_locking=False, name="HybridOptimizer"):
    super(HybridOptimizer, self).__init__(use_locking, name)
    self._optimizer_a = optimizer_a
    self._optimizer_b = optimizer_b

  def apply_gradients(self, grads_and_vars_a, grads_and_vars_b, global_step=None, name=None):
    op_a = self._optimizer_a.apply_gradients(grads_and_vars_a, name)
    op_b = self._optimizer_b.apply_gradients(grads_and_vars_b, name)
    return tf.group(*[op_a, op_b, global_step.assign_add(1)])

class SgdgOptimizer(optimizer.Optimizer):
  """Optimizer that implements stochastic gradient descent with momentum on G(1,n).

     References:
        - Minhyung Cho and Jaehyung Lee, Riemannian approach to batch normalization
          (https://arxiv.org/abs/1709.09603)
  """

  def __init__(self, learning_rate, momentum, grad_clip=None, use_locking=False, name="SgdgOptimizer"):
    super(SgdgOptimizer, self).__init__(use_locking, name)
    self._learning_rate = learning_rate
    self._momentum = momentum
    self._grad_clip = grad_clip

    # Tensor versions of the constructor arguments, created in _prepare().
    self._learning_rate_t = None
    self._momentum_t = None
    self._grad_clip_t = None

  def _prepare(self):
    self._learning_rate_t = tf.convert_to_tensor(self._learning_rate, name="learning_rate")
    self._momentum_t = tf.convert_to_tensor(self._momentum, name="momentum")
    if self._grad_clip!=None:
      self._grad_clip_t = tf.convert_to_tensor(self._grad_clip, name="grad_clip")
    else:
      self._grad_clip_t = None

  def _create_slots(self, var_list):
    for v in var_list:
      self._zeros_slot(v, "momentum", self._name)
          
  def _apply_dense(self, grad, var):
    mom = self.get_slot(var, "momentum")

    unity,_ = unit(var)     # for numerical stability
    h = gproj(unity, grad)

    if self._grad_clip_t!=None:
      h_hat = clip_by_norm(h, self._grad_clip_t)
    else:
      h_hat = h

    mom_new = self._momentum_t*mom - self._learning_rate_t*h_hat

    var_update = tf.assign(var, gexp(unity, mom_new))
    mom_update = tf.assign(mom, gpt(unity, mom_new))

    return tf.group(*[var_update, mom_update])

  def _apply_sparse(self, grad, var):
    raise NotImplementedError()


def reduce_shape(p):
  dim = len(p.get_shape())
  return [1]*(dim-1)+[p.get_shape().as_list()[-1]]

class AdamgOptimizer(optimizer.Optimizer):
  """Optimizer that implements Adam on G(1,n).

     References:
        - Minhyung Cho and Jaehyung Lee, Riemannian approach to batch normalization
          (https://arxiv.org/abs/1709.09603)
  """

  def __init__(self, learning_rate, beta1=0.9, beta2=0.99, epsilon=1e-8, grad_clip=None, use_locking=False, name="Adamg"):
    super(AdamgOptimizer, self).__init__(use_locking, name)
    self._lr = learning_rate
    self._beta1 = beta1
    self._beta2 = beta2
    self._epsilon = epsilon
    self._grad_clip = grad_clip

    # Tensor versions of the constructor arguments, created in _prepare().
    self._lr_t = None
    self._beta1_t = None
    self._beta2_t = None
    self._epsilon_t = None
    self._grad_clip_t = None

    # Variables to accumulate the powers of the beta parameters.
    # Created in _create_slots when we know the variables to optimize.
    self._beta1_power = None
    self._beta2_power = None

  def _get_beta_accumulators(self):
    return self._beta1_power, self._beta2_power

  def _create_slots(self, var_list):
    self._beta1_power = tf.Variable(self._beta1, name="beta1_power", trainable=False)
    self._beta2_power = tf.Variable(self._beta2, name="beta2_power", trainable=False)

    for v in var_list:
      dtype = v.dtype
      self._zeros_slot(v, "m", self._name)
      self._get_or_make_slot_with_initializer(v, tf.zeros_initializer(dtype), tf.TensorShape(reduce_shape(v)), dtype, "v", self._name)

  def _prepare(self):
    self._lr_t = tf.convert_to_tensor(self._lr, name="learning_rate")
    self._beta1_t = tf.convert_to_tensor(self._beta1, name="beta1")
    self._beta2_t = tf.convert_to_tensor(self._beta2, name="beta2")
    self._epsilon_t = tf.convert_to_tensor(self._epsilon, name="epsilon")
    if self._grad_clip!=None:
      self._grad_clip_t = tf.convert_to_tensor(self._grad_clip, name="grad_clip")
    else:
      self._grad_clip_t = None

  def _apply_dense(self, grad, var):
    m = self.get_slot(var, "m")
    v = self.get_slot(var, "v")

    unity,_ = unit(var)     # for numerical stability
    h = gproj(unity, grad)

    if self._grad_clip_t!=None:
      h_hat = clip_by_norm(h, self._grad_clip_t)
    else:
      h_hat = h


    mnew = self._beta1_t*m  + (1.0-self._beta1_t)*h_hat
    vnew = self._beta2_t*v  + (1.0-self._beta2_t)*xTy(h_hat,h_hat)

    alpha = tf.sqrt( 1- self._beta2_power) / (1.- self._beta1_power)
    deltas = (-alpha*self._lr_t) * mnew / tf.sqrt(vnew + self._epsilon_t)

    var_update = tf.assign(var, gexp(unity, deltas))
    m_update = tf.assign(m, gpt2(unity, mnew, deltas))
    v_update = tf.assign(v, vnew)


    return tf.group(*[var_update, m_update, v_update])

  def _apply_sparse(self, grad, var):
    raise NotImplementedError()


  def _finish(self, update_ops, name_scope):
    # Update the power accumulators.
    with tf.control_dependencies(update_ops):
      with ops.colocate_with(self._beta1_power):
        update_beta1 = self._beta1_power.assign(
            self._beta1_power * self._beta1_t,
            use_locking=self._use_locking)
        update_beta2 = self._beta2_power.assign(
            self._beta2_power * self._beta2_t,
            use_locking=self._use_locking)
    return tf.group(*update_ops + [update_beta1, update_beta2],
            name=name_scope)