Тензор потока: какие градиенты нужно было определить для пользовательской операции?

Question

Хотя есть много ссылок, показывающих, как зарегистрировать градиент, но мне все еще не очень ясно, какой именно градиент необходимо определить.

Некоторые похожие темы: Как зарегистрироватьПользовательский градиент для операции, состоящей из операций tf

Как определить только градиент для подграфа Tensorflow?

---

Хорошо, вот и мойвопрос:

У меня есть функция пересылки 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 ???, ???

Здесь говорит, что Результатом функции градиента должен быть список объектов 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 и 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.]]