Tensorflow Optimizer介绍

训练常规做法：

1. 定义损失函数loss

2. 使用optimizer进行minimize损失函数loss

minimize loss的由两部分组成：compute_gradients和apply_gradients

(1) compute_gradients

compute_gradients(
    loss,
    var_list=None,
    gate_gradients=GATE_OP,
    aggregation_method=None,
    colocate_gradients_with_ops=False,
    grad_loss=None
)
loss: 损失函数
var_list: 用于求梯度的变量
gate_gradients: 有三个可选项GATE_NONE，GATE_OP，GATE_GRAPH。GATE_NONE最高级别并发，代价是可能不可复现；GATE_OP在每个节点内部，计算完本节点全部梯度才更新参数，节点内部不并发，避免race condition；GATE_GRAPH最低级别并发，在计算好所有的梯度之后才更新参数
aggregation_method: 指定组合梯度的方法AggregationMethod
colocate_gradients_with_ops: 是否在多台gpu上并行计算梯度
grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`.

(2) apply_gradients

apply_gradients(
    grads_and_vars,
    global_step=None,
    name=None
)
grads_and_vars: List of (gradient, variable) pairs as returned by
  `compute_gradients()`.
global_step: Optional `Variable` to increment by one after the
  variables have been updated.
name: Optional name for the returned operation.  Default to the
  name passed to the `Optimizer` constructor.

以AdamOptimizer为例:
原论文的伪代码：

Tensorflow实现：
$lr_t := \text{learning\_rate} * \sqrt{(1 - \beta_2^t) / (1 - \beta_1^t)}$
$m_t := \beta_1 * m_{t-1} + (1 - \beta_1) * g$
$v_t := \beta_2 * v_{t-1} + (1 - \beta_2) * g * g$
$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$

代码见：https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/adam.py
稠密向量的操作：

  def _apply_dense(self, grad, var):
    m = self.get_slot(var, "m")# optiminzer额外参数加入_non_slot_dict
    v = self.get_slot(var, "v")
    beta1_power, beta2_power = self._get_beta_accumulators() # 获取t次方
    return training_ops.apply_adam(
        var,
        m,
        v,
        math_ops.cast(beta1_power, var.dtype.base_dtype),
        math_ops.cast(beta2_power, var.dtype.base_dtype),
        math_ops.cast(self._lr_t, var.dtype.base_dtype),
        math_ops.cast(self._beta1_t, var.dtype.base_dtype),
        math_ops.cast(self._beta2_t, var.dtype.base_dtype),
        math_ops.cast(self._epsilon_t, var.dtype.base_dtype),
        grad,
        use_locking=self._use_locking).op
  def _finish(self, update_ops, name_scope):
    # Update the power accumulators.
    with ops.control_dependencies(update_ops):
      beta1_power, beta2_power = self._get_beta_accumulators()
      with ops.colocate_with(beta1_power):
        update_beta1 = beta1_power.assign(
            beta1_power * self._beta1_t, use_locking=self._use_locking)
        update_beta2 = beta2_power.assign(
            beta2_power * self._beta2_t, use_locking=self._use_locking)
    return control_flow_ops.group(
        *update_ops + [update_beta1, update_beta2], name=name_scope)

其中beta1_power表示 $\beta_1^t$，beta2_power表示 $\beta_2^t$，t次方的实现在_finish函数中。
apply_adam具体实现代码见：https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/training_ops.cc

template <typename T>
struct ApplyAdam<CPUDevice, T> : ApplyAdamNonCuda<CPUDevice, T> {};

template <typename T>
struct ApplyAdamWithAmsgrad<CPUDevice, T> {
  void operator()(const CPUDevice& d, typename TTypes<T>::Flat var,
                  typename TTypes<T>::Flat m, typename TTypes<T>::Flat v,
                  typename TTypes<T>::Flat vhat,
                  typename TTypes<T>::ConstScalar beta1_power,
                  typename TTypes<T>::ConstScalar beta2_power,
                  typename TTypes<T>::ConstScalar lr,
                  typename TTypes<T>::ConstScalar beta1,
                  typename TTypes<T>::ConstScalar beta2,
                  typename TTypes<T>::ConstScalar epsilon,
                  typename TTypes<T>::ConstFlat grad) {
    const T alpha = lr() * Eigen::numext::sqrt(T(1) - beta2_power()) /
                    (T(1) - beta1_power());
    m.device(d) += (grad - m) * (T(1) - beta1()); # 更新m
    v.device(d) += (grad.square() - v) * (T(1) - beta2()); # 更新v
    vhat.device(d) = vhat.cwiseMax(v); 
    var.device(d) -= (m * alpha) / (vhat.sqrt() + epsilon()); # 更新参数
  }
};

稀疏向量的操作：

  def _apply_sparse_shared(self, grad, var, indices, scatter_add):
    beta1_power, beta2_power = self._get_beta_accumulators() # 获取t次方
    beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
    beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
    lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
    beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
    beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
    epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
    lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))
    # m_t = beta1 * m + (1 - beta1) * g_t
    m = self.get_slot(var, "m")
    m_scaled_g_values = grad * (1 - beta1_t)
    m_t = state_ops.assign(m, m * beta1_t, use_locking=self._use_locking)
    with ops.control_dependencies([m_t]):
      m_t = scatter_add(m, indices, m_scaled_g_values)
    # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
    v = self.get_slot(var, "v")
    v_scaled_g_values = (grad * grad) * (1 - beta2_t)
    v_t = state_ops.assign(v, v * beta2_t, use_locking=self._use_locking)
    with ops.control_dependencies([v_t]):
      v_t = scatter_add(v, indices, v_scaled_g_values)
    v_sqrt = math_ops.sqrt(v_t)
    var_update = state_ops.assign_sub(
        var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking)
    return control_flow_ops.group(*[var_update, m_t, v_t])
  def _finish(self, update_ops, name_scope):
    # Update the power accumulators.
    with ops.control_dependencies(update_ops):
      beta1_power, beta2_power = self._get_beta_accumulators()
      with ops.colocate_with(beta1_power):
        update_beta1 = beta1_power.assign(
            beta1_power * self._beta1_t, use_locking=self._use_locking)
        update_beta2 = beta2_power.assign(
            beta2_power * self._beta2_t, use_locking=self._use_locking)
    return control_flow_ops.group(
        *update_ops + [update_beta1, update_beta2], name=name_scope)

返回更新参数操作op的group（包括slot和non_slot)。

Bert实现AdamWeightDecayOptimizer的时候少了偏差纠正项，加入正则项。

  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
    """See base class."""
    assignments = []
    for (grad, param) in grads_and_vars:
      if grad is None or param is None:
        continue
      param_name = self._get_variable_name(param.name)
      m = tf.get_variable(
          name=param_name + "/adam_m",
          shape=param.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer())
      v = tf.get_variable(
          name=param_name + "/adam_v",
          shape=param.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer())
      # Standard Adam update.
      next_m = (
          tf.multiply(self.beta_1, m) + tf.multiply(1.0 - self.beta_1, grad))
      next_v = (
          tf.multiply(self.beta_2, v) + tf.multiply(1.0 - self.beta_2,
                                                    tf.square(grad)))
      update = next_m / (tf.sqrt(next_v) + self.epsilon)
      # Just adding the square of the weights to the loss function is *not*
      # the correct way of using L2 regularization/weight decay with Adam,
      # since that will interact with the m and v parameters in strange ways.
      #
      # Instead we want ot decay the weights in a manner that doesn't interact
      # with the m/v parameters. This is equivalent to adding the square
      # of the weights to the loss with plain (non-momentum) SGD.
      if self._do_use_weight_decay(param_name):
        update += self.weight_decay_rate * param
      update_with_lr = self.learning_rate * update
      next_param = param - update_with_lr
      assignments.extend(
          [param.assign(next_param),
           m.assign(next_m),
           v.assign(next_v)])
    return tf.group(*assignments, name=name)