tensorflow optimizer 源码阅读笔记

阅读数:18 2019 年 11 月 29 日 08:00

tensorflow optimizer源码阅读笔记

一直对 tf 中的自动求导机制比较好奇,它内部到底是怎么做梯度的反向传播的呢?所以最近阅读了 tensorflow/python/training/optimizer.py 的源码。其实 tf 的自动求导就是靠各式各样的 Optimizer 类进行的,我们只需要在程序中构建前向图,然后加上 Optimizer,再调用 minimize() 方法就可以完成梯度的反向传播。

Optimizer class 是所有 Optimizer 的基类(比如 GradientDescentOptimizer、AdamOptimizer 等),整个反向传播过程可分为三步,这三步仅需通过一个 minimize() 函数完成:

计算每一个部分的梯度,compute_gradients()
根据需要对梯度进行处理
把梯度更新到参数上,apply_gradients()

compute_gradients 函数

参数 gate_gradients:用于控制梯度计算过程的并行性

GATE_GRAPH 很好理解,即整个图中间的梯度计算(后向过程)和梯度更新是单独分开的,计算过程严格按照前向、后向、更新的步骤来,等到所有的参数都完成梯度计算之后,再统一发起更新。

GATE_NONE 和 GATE_OP 的差别在于梯度更新会不会影响到后续的其他计算。例如某个 op 有 n 个输入 x0,x1,…,xn−1,梯度的计算和更新需要对所有这 n 个输入求导,在 GATE_NONE 模式下,x0 的梯度计算完了之后,对 x0 的更新就马上开始了,那么在算其他输入(例如 xn−1)的梯度时,如果梯度项中含有 x0,就可能会出现“不可复现”的结果,因为每次算梯度时不一定哪一个梯度先算完呢。

GATE_OP 即产生一些控制依赖,确定某个变量不再会被用到之后才进行更新,保证正确性的同时最大化并行性。

核心代码:

复制代码
if var_list is None:
var_list = (
variables.trainable_variables() +
ops.get_collection(ops.GraphKeys.TRAINABLE_RESOURCE_VARIABLES))
else:
var_list = nest.flatten(var_list)
# pylint: disable=protected-access
var_list += ops.get_collection(ops.GraphKeys._STREAMING_MODEL_PORTS)
# pylint: enable=protected-access
processors = [_get_processor(v) for v in var_list]
if not var_list:
raise ValueError("No variables to optimize.")
var_refs = [p.target() for p in processors]
grads = gradients.gradients(
loss, var_refs, grad_ys=grad_loss,
gate_gradients=(gate_gradients == Optimizer.GATE_OP),
aggregation_method=aggregation_method,
colocate_gradients_with_ops=colocate_gradients_with_ops)
if gate_gradients == Optimizer.GATE_GRAPH:
grads = control_flow_ops.tuple(grads)
grads_and_vars = list(zip(grads, var_list))
self._assert_valid_dtypes(
[v for g, v in grads_and_vars
if g is not None and v.dtype != dtypes.resource])
return grads_and_vars

其中最核心的 gradients.gradients 函数,该函数可执行的功能为:根据原本计算图中所有的 op 创建一个顺序的 list,然后反向遍历这个 list,对每个需要求导并且能够求导的 op(即已经定义好了对应的梯度函数的 op)调用其梯度函数,然后沿着原本计算图的方向反向串起另一部分的计算图(输入输出互换,原本的数据 Tensor 换成梯度 Tensor)。

另外 _get_processor 函数可理解为一种快速更新 variables 的方法,每个 processor 都会包含一个 update_op 这样的函数来进行 variable 更新操作。

apply_gradients 函数

apply_gradients 函数根据前面求得的梯度,把梯度更新到参数上。

核心代码:

复制代码
converted_grads_and_vars = tuple(converted_grads_and_vars)
var_list = [v for g, v, _ in converted_grads_and_vars if g is not None]
if not var_list:
raise ValueError("No gradients provided for any variable: %s." %
([str(v) for _, v, _ in converted_grads_and_vars],))
with ops.init_scope():
self._create_slots(var_list)
update_ops = []
with ops.name_scope(name, self._name) as name:
self._prepare()
for grad, var, processor in converted_grads_and_vars:
if grad is None:
continue
# We colocate all ops created in _apply_dense or _apply_sparse
# on the same device as the variable.
# TODO(apassos): figure out how to get the variable name here.
if (context.executing_eagerly() or
resource_variable_ops.is_resource_variable(var)
and not var._in_graph_mode): # pylint: disable=protected-access
scope_name = ""
else:
scope_name = var.op.name
with ops.name_scope("update_" + scope_name), ops.colocate_with(var):
update_ops.append(processor.update_op(self, grad))
if global_step is None:
apply_updates = self._finish(update_ops, name)
else:
with ops.control_dependencies([self._finish(update_ops, "update")]):
with ops.colocate_with(global_step):
if isinstance(
global_step, resource_variable_ops.BaseResourceVariable):
# TODO(apassos): the implicit read in assign_add is slow; consider
# making it less so.
apply_updates = resource_variable_ops.assign_add_variable_op(
global_step.handle,
ops.convert_to_tensor(1, dtype=global_step.dtype),
name=name)
else:
apply_updates = state_ops.assign_add(global_step, 1, name=name)
if not context.executing_eagerly():
if isinstance(apply_updates, ops.Tensor):
apply_updates = apply_updates.op
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
if apply_updates not in train_op:
train_op.append(apply_updates)
return apply_updates

其中 self._create_slots 函数表示创建一些优化器自带的一些参数,比如 AdamOptimizer 的 m 和 v,[公式] 的 [公式] 次方(beta1_power)和 [公式] 的 [公式] 次方(beta2_power)。prepare() 函数的作用是在 apply 梯度前创建好所有必须的 tensors。

ops.colocate_with(var) 函数的作用是保证每个参数 var 的更新都在同一个 device 上。具体用法可以草考: https://stackoverflow.com/questions/45341067/what-is-colocate-with-used-for-in-tensorflow

ops.control_dependencies() 函数用来控制计算流图的,给图中的某些节点指定计算的顺序。代码中的意思就是先执行 update_ops 操作,然后再执行 global_step 的加 1 操作。

Optimizer 基类的这个方法为每个实现子类预留了 _create_slots(),_prepare(),_apply_dense(),_apply_sparse() 四个接口出来,后面新构建的 Optimizer 只需要重写或者扩展 Optimizer 类的某几个函数即可。

apply_gradients() 核心的部分就是对每个 variable 本身应用 assign,体现在 update_ops.append(processor.update_op(self, grad)),如果有 global_step 的话,global_step 需加个 1。

AdamOptimizer 类

位置: https://github.com/tensorflow/tensorflow/blob/818704e5cad0bed56f14281e9105eb39c1060bd1/tensorflow/python/training/adam.py

首先介绍一下 Adam 优化器的公式:

复制代码
def __init__(self,
learning_rate=0.001,
beta1=0.9,
beta2=0.999,
epsilon=1e-8,
use_locking=False,
name="Adam"):
super(AdamOptimizer, self).__init__(use_locking, name)
self._lr = learning_rate
self._beta1 = beta1
self._beta2 = beta2
self._epsilon = epsilon
# 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
def _prepare(self):
lr = self._call_if_callable(self._lr)
beta1 = self._call_if_callable(self._beta1)
beta2 = self._call_if_callable(self._beta2)
epsilon = self._call_if_callable(self._epsilon)
self._lr_t = ops.convert_to_tensor(lr, name="learning_rate")
self._beta1_t = ops.convert_to_tensor(beta1, name="beta1")
self._beta2_t = ops.convert_to_tensor(beta2, name="beta2")
self._epsilon_t = ops.convert_to_tensor(epsilon, name="epsilon")

上边 _init_ 函数可以看到,除了初始化时传进去的参数,优化器自身还存储了这些参数的 Tensor 版本,这些转换是在 _prepare 函数中通过 convert_to_tensor 方法来实现的。

复制代码
def _get_beta_accumulators(self):
with ops.init_scope():
if context.executing_eagerly():
graph = None
else:
graph = ops.get_default_graph()
return (self._get_non_slot_variable("beta1_power", graph=graph),
self._get_non_slot_variable("beta2_power", graph=graph))
def _create_slots(self, var_list):
# Create the beta1 and beta2 accumulators on the same device as the first
# variable. Sort the var_list to make sure this device is consistent across
# workers (these need to go on the same PS, otherwise some updates are
# silently ignored).
first_var = min(var_list, key=lambda x: x.name)
self._create_non_slot_variable(
initial_value=self._beta1, name="beta1_power", colocate_with=first_var)
self._create_non_slot_variable(
initial_value=self._beta2, name="beta2_power", colocate_with=first_var)
# Create slots for the first and second moments.
for v in var_list:
self._zeros_slot(v, "m", self._name)
self._zeros_slot(v, "v", self._name)

_create_slots 函数用来创建参数,被创建的参数有 [公式]、[公式]、[公式] 的 [公式] 次方(beta1_power)和 [公式] 的 [公式] 次方(beta2_power)。_get_beta_accumulators 函数是用来获取 [公式] 的 [公式] 次方(beta1_power)和 [公式] 的 [公式] 次方(beta2_power)的值。

复制代码
def _apply_dense(self, grad, var):
m = self.get_slot(var, "m")
v = self.get_slot(var, "v")
beta1_power, beta2_power = self._get_beta_accumulators()
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 _resource_apply_dense(self, grad, var):
m = self.get_slot(var, "m")
v = self.get_slot(var, "v")
beta1_power, beta2_power = self._get_beta_accumulators()
return training_ops.resource_apply_adam(
var.handle,
m.handle,
v.handle,
math_ops.cast(beta1_power, grad.dtype.base_dtype),
math_ops.cast(beta2_power, grad.dtype.base_dtype),
math_ops.cast(self._lr_t, grad.dtype.base_dtype),
math_ops.cast(self._beta1_t, grad.dtype.base_dtype),
math_ops.cast(self._beta2_t, grad.dtype.base_dtype),
math_ops.cast(self._epsilon_t, grad.dtype.base_dtype),
grad,
use_locking=self._use_locking)

函数 _apply_dense 和 _resource_apply_dense 的实现中分别使用了 training_ops.apply_adam 和 training_ops.resource_apply_adam 方法。具体实现于: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/training_ops.cc

复制代码
template <typename Device, typename T>
struct ApplyAdamNonCuda {
void operator()(const Device& d, typename TTypes<T>::Flat var,
typename TTypes<T>::Flat m, typename TTypes<T>::Flat v,
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, bool use_nesterov) {
// ...
T* var_ptr = var.data();
T* m_ptr = m.data();
T* v_ptr = v.data();
const T* g_ptr = grad.data();
const T alpha = lr() * Eigen::numext::sqrt(T(1) - beta2_power()) /
(T(1) - beta1_power());
if (use_nesterov) {
m += (g - m) * (T(1) - beta1());
v += (g.square() - v) * (T(1) - beta2());
var -= ((g * (T(1) - beta1()) + beta1() * m) * alpha) /
(v.sqrt() + epsilon());
} else {
m += (g - m) * (T(1) - beta1());
v += (g.square() - v) * (T(1) - beta2());
var -= (m * alpha) / (v.sqrt() + epsilon());
}
// ...
}
};

在上面的实现中,[公式] 和 [公式] 的更新公式和论文的形式上好像有些不同,但其实是一样的:

复制代码
def _apply_sparse(self, grad, var):
return self._apply_sparse_shared(grad.values, var, grad.indices,
lambda x, i, v: state_ops.scatter_add(x, i, v, use_locking=self._use_locking))
def _resource_apply_sparse(self, grad, var, indices):
return self._apply_sparse_shared(grad, var, indices, self._resource_scatter_add)
def _resource_scatter_add(self, x, i, v):
with ops.control_dependencies([resource_variable_ops.resource_scatter_add(x.handle, i, v)]):
return x.value()

函数 _apply_sparse 和 _resource_apply_sparse 主要用在稀疏向量的更新操作上,而具体的实现是在函数 _apply_sparse_shared 中。

复制代码
def _apply_sparse_shared(self, grad, var, indices, scatter_add):
beta1_power, beta2_power = self._get_beta_accumulators()
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])

scatter_add 函数作用是完成稀疏 Tensor 的加操作,其中代码中的参数 m 相当于 ref,indices 是索引,m_scaled_g_values 是更新的值。
tensorflow optimizer源码阅读笔记

那么现在分析下 _apply_sparse_shared 函数,首先获取所需要的参数值并存储到变量里,接着按照 Adam 算法的流程,首先计算学习率 [公式],接着计算两个 Momentum ,由于是稀疏 tensor 的更新,所以在算出更新值之后要使用 scatter_add 来完成加法操作, 最后将 var_update 和 m_t、v_t 的更新操作放进 control_flow_ops.group 中。

复制代码
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)和 [公式] 的 [公式] 次方(beta2_power)是在通过 _finish 函数计算的,通过之前存储的 [公式] 和 [公式] 的更新 op,并将这两个更新操作放进放到 control_flow_ops.group 中。 可以发现 adam 算法的所有的更新计算操作都会放进 control_flow_ops.group 中。

参考文献:

http://jcf94.com/2018/01/23/2018-01-23-tfunpacking2/
https://zhuanlan.zhihu.com/p/40870669
https://zhuanlan.zhihu.com/p/63500952
https://github.com/tensorflow/tensorflow/blob/818704e5cad0bed56f14281e9105eb39c1060bd1/tensorflow/python/training/optimizer.py
https://github.com/tensorflow/tensorflow/blob/818704e5cad0bed56f14281e9105eb39c1060bd1/tensorflow/python/training/adam.py#L145

本文转载自 Alex-zhai 知乎账号。

原文链接: https://zhuanlan.zhihu.com/p/87348147

评论

发布