python - 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 表示梯度函数的结果必须是表示每个输入的梯度的 Tensor 对象列表。

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

Answer 1

tf.gradients计算每个输出张量之和相对于输入张量中每个值的梯度。梯度操作接收您正在计算梯度的操作，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。

作为示例，我们以矩阵乘法(op 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 行的点积grad 和 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.]]