Как обеспечить пользовательский градиент в TensorFlow

Question

Я пытаюсь понять, как использовать функцию @tf.custom_gradient, доступную в TensorFlow 1.7 для предоставления собственного градиента вектора по отношению к вектору.Ниже приведен минимальный рабочий пример, который решает следующую проблему, чтобы получить dz/dx.

y = Ax
z = || y || ²

Кроме того, это прилагаемое изображение описывает решение, как и ожидалось, путем ручного вычисления

Если я не использую @tf.custom_gradient, то TensorFlow выдает желаемое решение, как и ожидалось.Мой вопрос заключается в том, как я могу предоставить собственный градиент для y = Ax?Мы знаем, что dy/dx = A^T, как показано в приведенном выше приложении, которое показывает шаги вычисления, которые соответствуют выводу TensorFlow.

import tensorflow as tf

#I want to write custom gradient for this function f1
def f1(A,x):
    y=tf.matmul(A,x,name='y')
    return y

#for y= Ax, the derivative is: dy/dx= transpose(A)
@tf.custom_gradient
def f2(A,x):
    y=f1(A,x)
    def grad(dzByDy): # dz/dy = 2y reaches here correctly.
        dzByDx=tf.matmul(A,dzByDy,transpose_a=True) 
        return dzByDx
    return y,grad


x= tf.constant([[1.],[0.]],name='x')
A= tf.constant([ [1., 2.], [3., 4.]],name='A')

y=f1(A,x) # This works as desired
#y=f2(A,x) #This line gives Error


z=tf.reduce_sum(y*y,name='z')

g=tf.gradients(ys=z,xs=x)

with tf.Session() as sess:
    print sess.run(g)

Answer 1

Поскольку ваша функция f2() имеет два входа, вы должны предоставить градиент для возврата к каждому из них.Вы видите ошибку:

Num градиентов 2, сгенерированных для имени операции: "IdentityN" [...] не соответствуют num входам 3

, по общему признанию, весьма загадочно,хоть.Предположим, что вы никогда не хотите рассчитывать d y / d A , вы можете просто вернуть None, dzByDx.Код ниже (проверено):

import tensorflow as tf

#I want to write custom gradient for this function f1
def f1(A,x):
    y=tf.matmul(A,x,name='y')
    return y

#for y= Ax, the derivative is: dy/dx= transpose(A)
@tf.custom_gradient
def f2(A,x):
    y=f1(A,x)
    def grad(dzByDy): # dz/dy = 2y reaches here correctly.
        dzByDx=tf.matmul(A,dzByDy,transpose_a=True) 
        return None, dzByDx
    return y,grad

x= tf.constant([[1.],[0.]],name='x')
A= tf.constant([ [1., 2.], [3., 4.]],name='A')

#y=f1(A,x) # This works as desired
y=f2(A,x) #This line gives Error

z=tf.reduce_sum(y*y,name='z')

g=tf.gradients(ys=z,xs=x)

with tf.Session() as sess:
    print sess.run( g )

вывод:

[массив ([[20.], [28.]], dtype = float32)]

по желанию.