Tensorflow：需要为自定义操作定义哪些梯度？

Question

虽然有很多参考资料展示了如何注册渐变，但我仍然不是很清楚到底需要定义什么样的渐变。

一些类似的主题：

好的，我的问题来了：

我有一个正向函数y = f(A,B)，其中每个函数的大小是：

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

假设我可以写出 y 的每个元素相对于 A 和 B 的每个元素的数学偏导数。dy/dA, dy/dB。我的问题是我应该return在渐变函数中做什么？

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

Here 说 梯度函数的结果必须是一个 Tensor 对象列表，代表每个输入的梯度。

y为标量，A、B为矩阵时定义的梯度非常容易理解。但是当y是矩阵，A，B也是矩阵时，那个梯度应该是多少呢？

Answer 1

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

以矩阵乘法(opMatMul)的梯度为例，你可以在tensorflow/python/ops/math_grad.py:

中找到

@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_a 和 transpose_b 都是 False 的情况，所以我们在第一个分支 if not t_a and not t_b: （也忽略 conj，这意味着复数）。 'a' 和 'b' 是这里的操作数，如前所述， grad 具有输出和相对于乘法结果中每个值的梯度。那么如果我将 a[0, 0] 增加一，事情会发生什么变化呢？基本上，乘积矩阵第一行中的每个元素都会增加 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] 将是 grad 的第 i 行和 b 的第 j 行的点积。您也可以对 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.]]

Tensorflow：需要为自定义操作定义哪些梯度？

Tensorflow: What gradients needed to be defined for custom operation?

python

gradient

keras

tensorflow