Home Articles

Tensorflow:需要为自定义操作定义哪些渐变?

Asked
Viewed 419 times
1

虽然有很多参考文献显示了如何注册渐变,但我仍然不太清楚需要定义什么样的渐变 .

一些类似的主题:How to register a custom gradient for a operation composed of tf operations

How Can I Define Only the Gradient for a Tensorflow Subgraph?


好的,这是我的问题:

我有一个前向函数 y = f(A,B) ,其中每个函数的大小为:

y: (batch_size, m, n)
A: (batch_size, a, a)
B: (batch_size, b, b)

enter image description here

假设我可以写下关于A和B的每个元素的y的每个元素的数学偏导数 . dy/dA, dy/dB . 我的问题是我应该在渐变函数中返回什么?

@ops.RegisterGradient("f")
def f_grad(op, grad):
    ...
    return ???, ???

Here表示渐变函数的结果必须是Tensor对象列表,表示相对于每个输入的渐变 .

y 是标量且 AB 是矩阵时,很容易理解要定义的梯度 . 但是当 y 是矩阵并且 AB 也是矩阵时,该梯度应该是什么?

1 Answer

  • 1

    tf.gradients计算每个输出张量之和相对于输入张量中每个值的梯度 . 渐变操作接收用于计算渐变的op, op 以及此时累积的渐变 grad . 在您的示例中, grad 将是与 y 具有相同形状的张量,并且每个值将是 y 中相应值的渐变 - 也就是说,如果 grad[0, 0] == 2 ,则意味着将 y[0, 0] 增加1将增加输出的总和张紧2(我知道,你可能已经清楚了) . 现在你必须为 AB 计算相同的东西 . 假设您认为将 A[2, 3] 增加1将使 y[0, 0] 增加3并且对 y 中的任何其他值都没有影响 . 这意味着将输出值的总和增加3×2 = 6,因此 A[2, 3] 的梯度将为6 .

    举个例子,让我们采用矩阵乘法的梯度(op MatMul ),你可以在_1844913中找到它:

    @ops.RegisterGradient("MatMul")
    def _MatMulGrad(op, grad):
      """Gradient for MatMul."""
    
      t_a = op.get_attr("transpose_a")
      t_b = op.get_attr("transpose_b")
      a = math_ops.conj(op.inputs[0])
      b = math_ops.conj(op.inputs[1])
      if not t_a and not t_b:
        grad_a = gen_math_ops.mat_mul(grad, b, transpose_b=True)
        grad_b = gen_math_ops.mat_mul(a, grad, transpose_a=True)
      elif not t_a and t_b:
        grad_a = gen_math_ops.mat_mul(grad, b)
        grad_b = gen_math_ops.mat_mul(grad, a, transpose_a=True)
      elif t_a and not t_b:
        grad_a = gen_math_ops.mat_mul(b, grad, transpose_b=True)
        grad_b = gen_math_ops.mat_mul(a, grad)
      elif t_a and t_b:
        grad_a = gen_math_ops.mat_mul(b, grad, transpose_a=True, transpose_b=True)
        grad_b = gen_math_ops.mat_mul(grad, a, transpose_a=True, transpose_b=True)
      return grad_a, grad_b
    

    我们将重点关注 transpose_atranspose_b 都是 False 的情况,因此我们在第一个分支 if not t_a and not t_b: (也忽略 conj ,这意味着复数值) . 'a'和'b'是这里的操作数,如前所述, grad 具有相对于乘法结果中每个值的输出之和的梯度 . 那么,如果我将_1844923增加一个,情况将如何变化?基本上,产品矩阵的第一行中的每个元素将增加 b 的第一行中的值 . 所以 a[0, 0] 的梯度是 b 的第一行和 grad 的第一行的点积 - 也就是说,我将增加每个输出值乘以每个输出值的累计梯度 . 如果你考虑一下,那么 grad_a = gen_math_ops.mat_mul(grad, b, transpose_b=True) 就是这么做的 . grad_a[0, 0] 将是 grad 的第一行和 b 的第一行的点积(因为我们在这里转置 b ),并且,通常, grad_a[i, j] 将是 gradi 行和 j 的点积 - 第 b 行 . 您也可以遵循 grad_b 的类似推理 .


    编辑:

    作为示例,请查看tf.gradients和注册的渐变如何相互关联:

    import tensorflow as tf
    # Import gradient registry to lookup gradient functions
    from tensorflow.python.framework.ops import _gradient_registry
    
    # Gradient function for matrix multiplication
    matmul_grad = _gradient_registry.lookup('MatMul')
    # A matrix multiplication
    a = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
    b = tf.constant([[6, 7, 8], [9, 10, 11]], dtype=tf.float32)
    c = tf.matmul(a, b)
    # Gradient of sum(c) wrt each element of a
    grad_c_a_1, = tf.gradients(c, a)
    # The same is obtained by backpropagating an all-ones matrix
    grad_c_a_2, _ = matmul_grad(c.op, tf.ones_like(c))
    # Multiply each element of c by itself, but stopping the gradients
    # This should scale the gradients by the values of c
    cc = c * tf.stop_gradient(c)
    # Regular gradients computation
    grad_cc_a_1, = tf.gradients(cc, a)
    # Gradients function called with c as backpropagated gradients
    grad_cc_a_2, _ = matmul_grad(c.op, c)
    with tf.Session() as sess:
        print('a:')
        print(sess.run(a))
        print('b:')
        print(sess.run(b))
        print('c = a * b:')
        print(sess.run(c))
        print('tf.gradients(c, a)[0]:')
        print(sess.run(grad_c_a_1))
        print('matmul_grad(c.op, tf.ones_like(c))[0]:')
        print(sess.run(grad_c_a_2))
        print('tf.gradients(c * tf.stop_gradient(c), a)[0]:')
        print(sess.run(grad_cc_a_1))
        print('matmul_grad(c.op, c)[0]:')
        print(sess.run(grad_cc_a_2))
    

    输出:

    a:
    [[1. 2.]
     [3. 4.]]
    b:
    [[ 6.  7.  8.]
     [ 9. 10. 11.]]
    c = a * b:
    [[24. 27. 30.]
     [54. 61. 68.]]
    tf.gradients(c, a)[0]:
    [[21. 30.]
     [21. 30.]]
    matmul_grad(c.op, tf.ones_like(c))[0]:
    [[21. 30.]
     [21. 30.]]
    tf.gradients(c * tf.stop_gradient(c), a)[0]:
    [[ 573.  816.]
     [1295. 1844.]]
    matmul_grad(c.op, c)[0]:
    [[ 573.  816.]
     [1295. 1844.]]
    

Related