通过py_function的高阶渐变

Question

我想知道如何通过tf2.0中的tf.py_function计算高阶梯度。以下示例（从tensorflow doc稍作修改）产生正确的dy_dx，而aa_x为None。谢谢。

import tensorflow as tf
import os

def huber(x, delta):
  if tf.abs(x) <= delta:
    return x*x/ (2*delta)
  else:
    return tf.abs(x)-delta/2.0



x = tf.constant ([2.0 ] )         
z = tf.constant ([1.0 ] )

with tf.GradientTape (persistent=True) as g0:
  g0.watch(x)

  with tf.GradientTape (persistent=True) as g :
    g.watch (x)
    y = tf.py_function(func=huber, inp=[x, 3.] , Tout=tf.float32  )

  dy_dx = g.gradient(y, x)
  aa = tf.reduce_sum(dy_dx *z )

aa_x = g0.gradient (aa, x)
print (dy_dx)
print (aa_x)

Answer 1

基于tf.py_function的文档，您无法计算更高的>1st导数。此函数允许将TensorFlow图中的计算表示为Python函数。特别是，它将Python函数func包装在once-differentiable TensorFlow操作中，并在启用了急切执行的情况下执行该函数。意味着您只能区分它[[once。

如果要获得

高阶导数

，则可以在Tensorflow 2.1.0中正常使用gradient函数。

修改的代码：

import tensorflow as tf     # Tensorflow 2.1.0
import os


def huber(x, delta):
  if tf.abs(x) <= delta:
    return x*x/ (2*delta)  ## x^2  / 2*delta  
                           ## x / delta   - 1st derivative
                           ## 1 / delta   - 2nd derivative
  else:
    return tf.abs(x)-delta/2.0 


x = tf.constant ([2.0 ])         
z = tf.constant ([1.0 ])

with tf.GradientTape (persistent=True) as g0:
  g0.watch(x)

  with tf.GradientTape (persistent=True) as g :
    g.watch (x)
    # y = tf.py_function(func=huber, inp=[x, 3.0] , Tout=tf.float32  )  # once-differentiable
    y= huber(x, 3.0)
  dy_dx = g.gradient(y, x)
  aa = tf.reduce_sum(dy_dx *z)

aa_x = g0.gradient(aa, x)

print (dy_dx)  #  tf.Tensor([0.6666667], shape=(1,), dtype=float32)
print (aa_x)   #  tf.Tensor([0.33333334], shape=(1,), dtype=float32)

您可以在此tf.py_wrap中阅读有关link功能的更多信息。