前段时间刚好写过一篇这方面的博客《自动微分》,最后介绍了一下TF自动求导的做法。具体内容贴在下面了
参考了知乎问题TensorFlow是如何求导的、StackOverflow问题Does tensorflow use automatic or symbolic gradients和TensorFlow关于eager execution模式的官方文档)
为了了解TensorFlow中自动微分的实现,需要先找到如何计算梯度。考虑到梯度常见的用处是最小化损失函数,因此可以先从损失函数如何优化的方向上探索,即从Optimizer
类的minimize
方法入手。这个方法调用了compute_gradients
方法以获得参数的梯度(然后会调用apply_gradients
以利用梯度更新参数,与本文讨论的内容暂时没有什么关系,所以先略去了)。由于TensorFlow有两种计算梯度的方法:一种是经典的静态图模式,一种是新加入的动态图模式(官方说法是eager execution模式),因此对于不同模式,compute_gradients
采取了不同的实现逻辑
TensorFlow的经典模式是先建立一个静态图,然后这个静态图在一个会话里执行。在这种模式下,compute_gradients
方法进一步调用tensorflow.python.ops.gradients_impl
里的gradients
方法
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)
其中loss
是计算损失值的张量,var_refs
是变量列表,grad_ys
存储计算出的梯度,gate_gradients
是一个布尔变量,指示所有梯度是否在使用前被算出,如果设为True
,可以避免竞争条件。不过gradients
方法在实现上用途更广泛一些,简单说,它就是为了计算一组输出张量ys = [y0, y1, ...]
对输入张量xs = [x0, x1, ...]
的梯度,对每个xi
有grad_i = sum[dy_j/dx_i for y_j in ys]
。默认情况下,grad_loss
是None
,此时grad_ys
被初始化为全1向量
gradients
实际上直接调用内部方法_GradientsHelper
@tf_export("gradients") def gradients(ys, xs, grad_ys=None, name="gradients", colocate_gradients_with_ops=False, gate_gradients=False, aggregation_method=None, stop_gradients=None): # Creating the gradient graph for control flow mutates Operations. # _mutation_lock ensures a Session.run call cannot occur between creating and # mutating new ops. with ops.get_default_graph()._mutation_lock(): # pylint: disable=protected-access return _GradientsHelper(ys, xs, grad_ys, name, colocate_gradients_with_ops, gate_gradients, aggregation_method, stop_gradients)
这个方法会维护两个重要变量
queue
,队列里存放计算图里所有出度为0的操作符grads
,字典的键是操作符本身,值是该操作符每个输出端收到的梯度列表反向传播求梯度时,每从队列中弹出一个操作符,都会把它输出变量的梯度加起来(对应全微分定理)得到out_grads
,然后获取对应的梯度计算函数grad_fn
。操作符op
本身和out_grads
会传递给grad_fn
做参数,求出输入的梯度
if grad_fn: # If grad_fn was found, do not use SymbolicGradient even for # functions. in_grads = _MaybeCompile(grad_scope, op, func_call, lambda: grad_fn(op, *out_grads)) else: # For function call ops, we add a 'SymbolicGradient' # node to the graph to compute gradients. in_grads = _MaybeCompile(grad_scope, op, func_call, lambda: _SymGrad(op, out_grads, xs))
(不过这里似乎说明TensorFlow是自动微分和符号微分混用的)
该操作符处理以后,会更新所有未经过处理的操作符的出度和queue
(实际上就是一个拓扑排序的过程)。这样,当queue
为空的时候,整个计算图处理完毕,可以得到每个参数的梯度
静态图模式下梯度计算的调用过程大致如下所示
Optimizer.minimize |---Optimizer.compute_gradients |---gradients (gradients_impl.py) |---_GradientsHelper (gradients_impl.py)
前面提到,在_GradientsHelper
函数里要调用一个grad_fn
函数,该函数用来计算给定操作符的梯度。在TensorFlow里,每个计算图都可以分解到操作符(op)层级,每个操作符都会定义一个对应的梯度计算函数。例如,在python/ops/math_grad.py里定义的Log函数的梯度
@ops.RegisterGradient("Log") def _LogGrad(op, grad): """Returns grad * (1/x).""" x = op.inputs[0] with ops.control_dependencies([grad]): x = math_ops.conj(x) return grad * math_ops.reciprocal(x)
返回的就是已有梯度和x
倒数的积,对应于
注意每个函数都使用了装饰器RegisterGradient
包装,对有m个输入,n个输出的操作符,相应的梯度函数需要传入两个参数
返回m个张量对象,代表对每个输入的梯度
大部分操作符的梯度计算方式已经由框架给出,但是也可以自定义操作和对应的梯度计算函数。假设要定义一个Sub
操作,接受两个输入x
和y
,输出一个x-y
,那么这个函数是
显然有
那么对应的代码就是
@tf.RegisterGradient("Sub") def _sub_grad(unused_op, grad): return grad, tf.negative(grad)
在动态图模式下,TensorFlow不需要预先定义好完整的计算图,每个操作也可以返回具体的值,方便调试。下面给出了一个使用动态图求解线性回归的例子(改动自官方示例代码)
import tensorflow as tf tf.enable_eager_execution() NUM_EXAMPLES = 1000 training_inputs = tf.random_normal([NUM_EXAMPLES]) noise = tf.random_normal([NUM_EXAMPLES]) training_outputs = training_inputs * 3 + 2 + noise def prediction(x, w, b): return x * w + b # A loss function using mean-squared error def loss(weights, biases): error = prediction(training_inputs, weights, biases) - training_outputs return tf.reduce_mean(tf.square(error)) train_steps = 200 learning_rate = 0.1 # Start with arbitrary values for W and B on the same batch of data weight = tf.Variable(5.) bias = tf.Variable(10.) optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate) for i in range(20): print("Initial loss: {:.3f}".format(loss(weight, bias))) optimizer.minimize(lambda: loss(weight, bias)) print("Final loss: {:.3f}".format(loss(weight, bias))) print("W = {}, B = {}".format(weight.numpy(), bias.numpy()))
仍然以Optimizer
类的minimize
方法为入口,跟进到compute_gradients
方法,可以看到在动态图模式下,相关代码比较简短
if callable(loss): with backprop.GradientTape() as tape: if var_list is not None: tape.watch(var_list) loss_value = loss() # Scale loss if using a "mean" loss reduction and multiple towers. # Have to be careful to call distribute_lib.get_loss_reduction() # *after* loss() is evaluated, so we know what loss reduction it uses. # TODO(josh11b): Test that we handle weight decay in a reasonable way. if (distribute_lib.get_loss_reduction() == variable_scope.VariableAggregation.MEAN): num_towers = distribution_strategy_context.get_distribution_strategy( ).num_towers if num_towers > 1: loss_value *= (1. / num_towers) if var_list is None: var_list = tape.watched_variables() grads = tape.gradient(loss_value, var_list, grad_loss) return list(zip(grads, var_list))
之前看到过一个比喻:自动微分的工作原理就像是录制一盘磁带:前向计算所有操作的时候,实际上是在录制正在进行的操作。等到录制结束,倒带播放,就得到了梯度。TensorFlow也遵循了这样的比喻,所以在动态图模式下自动微分的灵魂是一个GradientTape
(“磁带”)类的对象,通过这个对象记录数据,求出梯度
在该方法的第一步里,GradientTape
类对象tape
会在自己的context下“观察”所有需要被记录的对象。默认情况下,使用tf.Variable
或tf.get_variable()
创建的对象都是trainable
的,也是会被观察的(自动放在watched_variables
里)。然后,调用gradient
方法来计算所有被观察对象的梯度,核心代码为
flat_grad = imperative_grad.imperative_grad( _default_vspace, self._tape, nest.flatten(target), flat_sources, output_gradients=output_gradients)
这个函数最后会调用一个C++实现的ComputeGradient
函数,其伪代码大致如下
template <typename Gradient, typename BackwardFunction, typename TapeTensor> // 使用了传统C的约定,返回一个状态码,结果保存在result变量里 // 核心思想还是对有向图使用拓扑排序,找到出度为0的点,聚合上游梯度,求出下游梯度 Status GradientTape<Gradient, BackwardFunction, TapeTensor>::ComputeGradient( const VSpace<Gradient, BackwardFunction, TapeTensor>& vspace, gtl::ArraySlice<int64> target_tensor_ids, gtl::ArraySlice<int64> source_tensor_ids, gtl::ArraySlice<Gradient*> output_gradients, std::vector<Gradient*>* result) { // 构建一个输入张量的集合 gtl::FlatSet<int64> sources_set(source_tensor_ids.begin(), source_tensor_ids.end()); // 初始化,找到所有与输出张量有关的op,计算它们的出度、引用数等 BackpropInitialState<BackwardFunction, TapeTensor> state = PrepareBackprop( target_tensor_ids, tensor_tape_, &op_tape_, sources_set, persistent_); // 找到所有出度为0的op std::vector<int64> op_stack = InitialStack(state.op_tape, state.op_missing_tensor); gtl::FlatMap<int64, std::vector<Gradient*>> gradients; // 将所有最终输出的输出梯度设为1 Status s = InitialGradients(vspace, target_tensor_ids, output_gradients, tensor_tape_, state.op_tape, &gradients); while (!op_stack.empty()) { 获得一个op,从state.op_tape擦除之 获取输出的梯度(上游梯度) 计算输入的梯度(下游梯度)。大部分操作是使用CallBackwardFunction来完成 对每个输入张量,看它是哪些op的输出张量,将该op“未计算梯度的输出张量”的计数减1。当该计数降为0时,这个op相当于出度为0,可以放入op_stack } 聚合所有源向量的梯度 }
可以看出核心计算梯度的方法是调用CallBackwardFunction
。这个方法调用了操作符对应的反向传播函数backward_function
,而操作符和反向传播函数的对应关系会在“录制磁带”时记录
(这里有一个疑点,怀疑TensorFlow是如下逻辑:
只是猜想,不是很确定,欢迎证明/证伪)
动态计算图下梯度计算的调用过程大致如下所示
Optimizer.minimize |---Optimizer.compute_gradients |---GradientTape.gradient |---imperative_grad |---TFE_Py_TapeGradient (python/eager/pywrap_tfe_src.cc) |---GradientTape<>::ComputeGradient (c/eager/tape.h)