Python tf.Graph.gradient_override_map用法及代码示例

用法

@tf_contextlib.contextmanager
gradient_override_map(
    op_type_map
)

实验：用于覆盖梯度函数的上下文管理器。

此上下文管理器可用于覆盖将用于上下文范围内操作的梯度函数。

例如：

@tf.RegisterGradient("CustomSquare")
def _custom_square_grad(op, grad):
  # ...

with tf.Graph().as_default() as g:
  c = tf.constant(5.0)
  s_1 = tf.square(c)  # Uses the default gradient for tf.square.
  with g.gradient_override_map({"Square":"CustomSquare"}):
    s_2 = tf.square(s_2)  # Uses _custom_square_grad to compute the
                          # gradient of s_2.