如何在TensorFlow中提供自定义渐变

时间:2018-04-25 19:32:32

标签: python-2.7 tensorflow

我试图了解如何使用TensorFlow 1.7中提供的@tf.custom_gradient函数来提供矢量相对于矢量的自定义渐变。下面的代码是解决以下问题的最小工作示例,以获取dz/dx

Y =斧
z = || y || 2

Also, this attached image describes the solution as expected by manually calulation

如果我不使用@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)

1 个答案:

答案 0 :(得分:0)

由于您的函数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)]

根据需要。