基于TensorFlow网站上的示例引用： https://www.tensorflow.org/api_docs/python/tf/custom_gradient
@tf.custom_gradientdef op_with_fused_backprop(x): y, x_grad = fused_op(x) def first_order_gradient(dy): @tf.custom_gradient def first_order_custom(unused_x): def second_order_and_transpose(ddy): return second_order_for_x(...), gradient_wrt_dy(...) return x_grad, second_order_and_transpose return dy * first_order_custom(x) return y, first_order_gradient关于second_order_and_transpose(ddy)为何返回两个对象的细节不够充分。根据tf.custom_gradient的文档，grad_fn（即second_order_and_transpose()）应该返回一个张量列表，这些张量是unused_x对dy的导数。甚至不清楚为什么他们将其命名为unused_x。有人对这个例子或一般情况下为高阶导数创建自定义梯度有任何想法吗？
回答：
-
关于
second_order_and_transpose(ddy)为何返回两个对象的细节不够充分。
根据我对一些例子的尝试，我认为你是对的。官方文档有些模糊（或不正确）。second_order_and_transpose(ddy)应该只返回一个对象，即计算出的二阶梯度。
-
甚至不清楚为什么他们将其命名为
unused_x。
这是个棘手的部分。unused_x解释了他们为何这样命名（因为你永远不会使用它…）。这里的目标是将你的二阶计算函数包装在一个名为first_order_custom的函数中。你从fused_op计算x的梯度，并使用它作为返回值，而不是unused_x。
为了更清楚地说明这一点，我传递了一个从官方文档扩展的示例，以定义log1pexp的二阶梯度：
注意： 二阶梯度在数值上不稳定，所以我们使用(1 - tf.exp(x))来表示它，只是为了让我们的生活更轻松。
@tf.custom_gradientdef log1pexp2(x): e = tf.exp(x) y = tf.math.log(1 + e) x_grad = 1 - 1 / (1 + e) def first_order_gradient(dy): @tf.custom_gradient def first_order_custom(unused_x): def second_order_gradient(ddy): # 我们定义二阶梯度为(1 - e) return ddy * (1 - e) return x_grad, second_order_gradient return dy * first_order_custom(x) return y, first_order_gradient要测试脚本，只需运行：
import tensorflow as tf@tf.custom_gradientdef log1pexp2(x): e = tf.exp(x) y = tf.math.log(1 + e) x_grad = 1 - 1 / (1 + e) def first_order_gradient(dy): @tf.custom_gradient def first_order_custom(unused_x): def second_order_gradient(ddy): # 我们定义二阶梯度为(1 - e) return ddy * (1 - e) return x_grad, second_order_gradient return dy * first_order_custom(x) return y, first_order_gradientx1 = tf.constant(1.)y1 = log1pexp2(x1)dy1 = tf.gradients(y1, x1)ddy1 = tf.gradients(dy1, x1)x2 = tf.constant(100.)y2 = log1pexp2(x2)dy2 = tf.gradients(y2, x2)ddy2 = tf.gradients(dy2, x2)with tf.Session() as sess: print('x=1, dy1:', dy1[0].eval(session=sess)) print('x=1, ddy1:', ddy1[0].eval(session=sess)) print('x=100, dy2:', dy2[0].eval(session=sess)) print('x=100, ddy2:', ddy2[0].eval(session=sess))结果：
x=1, dy1: 0.7310586x=1, ddy1: -1.7182817x=100, dy2: 1.0x=100, ddy2: -inf