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.