Разбиение вычисления градиента TensorFlow на две (или более) части

Question

Можно ли использовать TensorFlow's tf.gradients() функция по частям, то есть - вычислить градиент от потери относительно некоторого тензора и этого тензора относительно веса, а затем умножить их, чтобы получить исходный градиент от потери к весу?

Например, пусть W,b будь немного весов, пусть x быть входом в сеть, и пусть y0 обозначать этикетки.

Предположим, что прямой граф

h=Wx+b
y=tanh(h)
loss=mse(y-y0)

Мы можем рассчитать tf.gradients(loss,W) а затем применить (пропуская некоторые детали) optimizer.apply_gradients() обновлять W,

Затем я пытаюсь извлечь промежуточный тензор, используя var=tf.get_default_graph().get_tensor_by_name(...), а затем рассчитать два градиента: g1=tf.gradients(loss,var) а также g2=tf.gradients(var,W), Тогда я бы, по правилу цепочки, ожидал g1 а также g2 работать так, чтобы я мог написать g=g1*g2 в каком-то смысле и вернусь tf.gradients(loss,W),

К сожалению, это не случай. Размеры неверны. Размеры каждого градиента будут такими же, как у "переменной wrt", поэтому между первым градиентом и вторым не будет соответствия. Что мне не хватает, и как я могу это сделать?

Благодарю.

Answer 1

tf.gradients будет суммировать по градиентам входного тензора. Чтобы этого избежать, нужно разбить тензор на скаляры и применить tf.gradients каждому из них:

import tensorflow as tf

x = tf.ones([1, 10])

w = tf.get_variable("w", initializer=tf.constant(0.5, shape=[10, 5]))
out = tf.matmul(x, w)
out_target = tf.constant(0., shape=[5])

loss = tf.reduce_mean(tf.square(out - out_target))

grad = tf.gradients(loss, x)

part_grad_1 = tf.gradients(loss, out)
part_grad_2 = tf.concat([tf.gradients(i, x) for i in tf.split(out, 5, axis=1)], axis=1)

grad_by_parts = tf.matmul(part_grad_1, part_grad_2)

init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init)
    print(sess.run([grad]))
    print(sess.run([grad_by_parts]))

Answer 2

Из документов, tf.gradients (акцент мой)

строит символические производные суммы ys по x в xs.

Если какой-либо тензор в ys в многомерном reduce_sumЗначение med перед результирующим списком скаляров суммируется перед дифференцированием. Вот почему выходной градиент имеет тот же размер, что и xs,

Это также объясняет, почему потери могут быть многомерными в тензорном потоке: они неявно суммируются перед дифференцированием.

Answer 3

для будущих читателей:

Tensorflow добился некоторых успехов, и что касается tf2.7 (и, возможно, даже более ранних версий) , вы можете использовать tf.GradientTape.jacobian, чтобы избежать суммирования размеров цели .

https://www.tensorflow.org/guide/advanced_autodiff#jacobians