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

Question

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

一些类似的话题： 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)

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

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

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

当y是标量而A、B是矩阵时，很容易理解要定义的梯度。但是当y是矩阵，A、B也是矩阵时，那个梯度应该是什么？

Answer 1

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

我们以矩阵乘法的梯度为例（操作MatMul），你可以在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] 将是i-th 行grad 和j-th 行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.]]