<gh_stars>0
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""PowerSign oprimizer"""

from __future__ import absolute_import, division, print_function
from tensorflow.python.ops import math_ops, state_ops
from tensorflow.python.training import optimizer

class PowerSignOptimizer(optimizer.Optimizer):
  def __init__(self, alpha=None, decay=0.9, lr=1e-4,
               use_locking=False, name="PowerSign"):
    """Construct a new PowerSignOptimizer.

    https://arxiv.org/abs/1709.07417

    Args:
      alpha: Float. The base of PowerSign update.
      decay: Float. Decay to use to maintain the moving averages
                    of trained variables.
      lr: Float. Learning rate.
      use_locking: If True use locks for update operations.
      name: Optional name for the operations created when applying gradients.
        Defaults to "PowerSign".
      """
    super(PowerSignOptimizer, self).__init__(use_locking, name)
    # self._ema = moving_averages.ExponentialMovingAverage(decay, num_updates=num_updates)
    self._variable_map = None
    self._alpha = alpha
    self._lr = lr
    self._decay = decay
    self._name = "power_sign_optimizer"

  def _create_slots(self, var_list):
    for v in var_list:
      self._zeros_slot(v, "m", self._name)

  # def apply_gradients(self, grads_and_vars, global_step=None, name=None):
  #  grad_list = [x[0] for x in grads_and_vars if x[0] is not None]
  #  ma_op = self._ema.apply(grad_list)

  #  train_op = super().apply_gradients(grads_and_vars, global_step=global_step, name=name)
  #  return control_flow_ops.group(train_op, ma_op, name="train_with_avg")

  def _apply_dense(self, grad, var):
    return self._apply_power_sign(grad, var)

  def _resource_apply_dense(self, grad, var):
    return self._apply_power_sign(grad.handle, var.handle)

  def _apply_sparse(self, grad, var):
    return self._apply_power_sign(grad, var)

  def _apply_power_sign(self, grad, var):
    m = self.get_slot(var, "m")
    lr = math_ops.cast(self._lr, var.dtype.base_dtype),
    decay = math_ops.cast(self._decay, var.dtype.base_dtype)
    
    m = (grad - m) * decay + m 
    same_sign = math_ops.sign(m) * math_ops.sign(grad)

    if self._alpha is None:
      delta = math_ops.exp(same_sign) * grad
    else:
      alpha = math_ops.cast(self._alpha, var.dtype.base_dtype),
      delta = math_ops.pow(alpha, same_sign) * grad
    var = state_ops.assign_sub(var, delta * lr, use_locking=True)

    return var


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""PowerSign oprimizer"""

from __future__ import absolute_import, division, print_function
from tensorflow.python.ops import math_ops, state_ops
from tensorflow.python.training import optimizer

class PowerSignOptimizer(optimizer.Optimizer):
  def __init__(self, alpha=None, decay=0.9, lr=1e-4,
               use_locking=False, name="PowerSign"):
    """Construct a new PowerSignOptimizer.

    https://arxiv.org/abs/1709.07417

    Args:
      alpha: Float. The base of PowerSign update.
      decay: Float. Decay to use to maintain the moving averages
                    of trained variables.
      lr: Float. Learning rate.
      use_locking: If True use locks for update operations.
      name: Optional name for the operations created when applying gradients.
        Defaults to "PowerSign".
      """
    super(PowerSignOptimizer, self).__init__(use_locking, name)
    # self._ema = moving_averages.ExponentialMovingAverage(decay, num_updates=num_updates)
    self._variable_map = None
    self._alpha = alpha
    self._lr = lr
    self._decay = decay
    self._name = "power_sign_optimizer"

  def _create_slots(self, var_list):
    for v in var_list:
      self._zeros_slot(v, "m", self._name)

  # def apply_gradients(self, grads_and_vars, global_step=None, name=None):
  #  grad_list = [x[0] for x in grads_and_vars if x[0] is not None]
  #  ma_op = self._ema.apply(grad_list)

  #  train_op = super().apply_gradients(grads_and_vars, global_step=global_step, name=name)
  #  return control_flow_ops.group(train_op, ma_op, name="train_with_avg")

  def _apply_dense(self, grad, var):
    return self._apply_power_sign(grad, var)

  def _resource_apply_dense(self, grad, var):
    return self._apply_power_sign(grad.handle, var.handle)

  def _apply_sparse(self, grad, var):
    return self._apply_power_sign(grad, var)

  def _apply_power_sign(self, grad, var):
    m = self.get_slot(var, "m")
    lr = math_ops.cast(self._lr, var.dtype.base_dtype),
    decay = math_ops.cast(self._decay, var.dtype.base_dtype)
    
    m = (grad - m) * decay + m 
    same_sign = math_ops.sign(m) * math_ops.sign(grad)

    if self._alpha is None:
      delta = math_ops.exp(same_sign) * grad
    else:
      alpha = math_ops.cast(self._alpha, var.dtype.base_dtype),
      delta = math_ops.pow(alpha, same_sign) * grad
    var = state_ops.assign_sub(var, delta * lr, use_locking=True)

    return var
# Copyright 2016 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== PowerSign oprimizer Construct a new PowerSignOptimizer.      https://arxiv.org/abs/1709.07417      Args:       alpha: Float. The base of PowerSign update.       decay: Float. Decay to use to maintain the moving averages                     of trained variables.       lr: Float. Learning rate.       use_locking: If True use locks for update operations.       name: Optional name for the operations created when applying gradients.         Defaults to "PowerSign". # self._ema = moving_averages.ExponentialMovingAverage(decay, num_updates=num_updates) # def apply_gradients(self, grads_and_vars, global_step=None, name=None): #  grad_list = [x[0] for x in grads_and_vars if x[0] is not None] #  ma_op = self._ema.apply(grad_list) #  train_op = super().apply_gradients(grads_and_vars, global_step=global_step, name=name) #  return control_flow_ops.group(train_op, ma_op, name="train_with_avg")