Python tf.custom_gradient用法及代碼示例

使用自定義漸變定義函數的裝飾器。

用法

tf.custom_gradient(
    f=None
)

這個裝飾器允許對操作序列的梯度進行細粒度控製。由於多種原因，這可能很有用，包括為一係列操作提供更有效或數值穩定的梯度。

例如，考慮在交叉熵和對數似然計算中經常出現的以下函數：

def log1pexp(x):
  return tf.math.log(1 + tf.exp(x))

由於數值不穩定性，此函數在 x=100 處評估的梯度為 NaN。例如：

x = tf.constant(100.)
y = log1pexp(x)
dy_dx = tf.gradients(y, x) # Will be NaN when evaluated.

梯度表達式可以解析簡化以提供數值穩定性：

@tf.custom_gradient
def log1pexp(x):
  e = tf.exp(x)
  def grad(upstream):
    return upstream * (1 - 1 / (1 + e))
  return tf.math.log(1 + e), grad

使用此定義，x = 100 處的梯度 dy_dx 將被正確評估為 1.0。

變量upstream 被定義為上遊梯度。即來自該層的所有層或函數的梯度。上麵的例子沒有上遊函數，因此 upstream = dy/dy = 1.0 。

假設 x_i 是前向傳遞中的 log1pexp x_1 = x_1(x_0) , x_2 = x_2(x_1) , ..., x_i = x_i(x_i-1) , ..., x_n = x_n(x_n-1) 。通過鏈式法則，我們知道 dx_n/dx_0 = dx_n/dx_n-1 * dx_n-1/dx_n-2 * ... * dx_i/dx_i-1 * ... * dx_1/dx_0 。

在這種情況下，我們當前函數的梯度定義為 dx_i/dx_i-1 = (1 - 1 / (1 + e)) 。上遊梯度 upstream 將是 dx_n/dx_n-1 * dx_n-1/dx_n-2 * ... * dx_i+1/dx_i 。上遊梯度乘以當前梯度然後向下遊傳遞。

如果函數將多個變量作為輸入，grad 函數也必須返回相同數量的變量。我們以函數z = x * y為例。

@tf.custom_gradient
def bar(x, y):
  def grad(upstream):
    dz_dx = y
    dz_dy = x
    return upstream * dz_dx, upstream * dz_dy
  z = x * y
  return z, grad
x = tf.constant(2.0, dtype=tf.float32)
y = tf.constant(3.0, dtype=tf.float32)
with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  tape.watch(y)
  z = bar(x, y)
z
<tf.Tensor:shape=(), dtype=float32, numpy=6.0>
tape.gradient(z, x)
<tf.Tensor:shape=(), dtype=float32, numpy=3.0>
tape.gradient(z, y)
<tf.Tensor:shape=(), dtype=float32, numpy=2.0>

嵌套自定義漸變可能會導致不直觀的結果。默認行為與n-th 階導數不對應。例如

@tf.custom_gradient
def op(x):
  y = op1(x)
  @tf.custom_gradient
  def grad_fn(dy):
    gdy = op2(x, y, dy)
    def grad_grad_fn(ddy): # Not the 2nd order gradient of op w.r.t. x.
      return op3(x, y, dy, ddy)
    return gdy, grad_grad_fn
  return y, grad_fn

grad_grad_fn函數將計算grad_fn相對於dy的一階梯度，用於從backward-mode梯度圖生成forward-mode梯度圖，但與二階梯度不同op 相對於 x 。

相反，將嵌套的 @tf.custom_gradients 包裝在另一個函數中：

@tf.custom_gradient
def 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

內部 @tf.custom_gradient 裝飾函數的附加參數控製最內部函數的預期返回值。

上麵的示例說明了如何為不從變量讀取的函數指定自定義漸變。以下示例使用變量，這些變量需要特殊處理，因為它們是 forward 函數的有效輸入。

weights = tf.Variable(tf.ones([2]))  # Trainable variable weights
@tf.custom_gradient
def linear_poly(x):
  # Creating polynomial
  poly = weights[1] * x + weights[0]

  def grad_fn(dpoly, variables):
    # dy/dx = weights[1] and we need to left multiply dpoly
    grad_xs = dpoly * weights[1]  # Scalar gradient

    grad_vars = []  # To store gradients of passed variables
    assert variables is not None
    assert len(variables) == 1
    assert variables[0] is weights
    # Manually computing dy/dweights
    dy_dw = dpoly * tf.stack([x ** 1, x ** 0])
    grad_vars.append(
        tf.reduce_sum(tf.reshape(dy_dw, [2, -1]), axis=1)
    )
    return grad_xs, grad_vars
  return poly, grad_fn
x = tf.constant([1., 2., 3.])
with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  poly = linear_poly(x)
poly # poly = x + 1
<tf.Tensor:shape=(3,),
  dtype=float32,
  numpy=array([2., 3., 4.], dtype=float32)>
tape.gradient(poly, x)  # conventional scalar gradient dy/dx
<tf.Tensor:shape=(3,),
  dtype=float32,
  numpy=array([1., 1., 1.], dtype=float32)>
tape.gradient(poly, weights)
<tf.Tensor:shape=(2,), dtype=float32, numpy=array([6., 3.], dtype=float32)>

上麵的示例說明了可訓練變量 weights 的用法。在示例中，內部 grad_fn 接受額外的 variables 輸入參數，並返回額外的 grad_vars 輸出。如果 forward 函數讀取任何變量，則傳遞該額外參數。您需要計算梯度 w.r.t。每個 variables 並將其輸出為 grad_vars 的列表。請注意，當 forward 函數中沒有使用任何變量時，variables 的默認值設置為 None。

應該注意 tf.GradientTape 仍在觀察 tf.custom_gradient 的前向傳遞，並將使用它觀察的操作。因此，在磁帶仍在觀看時調用 tf.function 會導致構建梯度圖。如果在 tf.function 中使用了沒有注冊梯度的操作，則會引發 LookupError。

用戶可以插入tf.stop_gradient 來自定義此行為。這在下麵的示例中得到了證明。 tf.random.shuffle 沒有注冊的漸變。結果 tf.stop_gradient 用於避免 LookupError 。

x = tf.constant([0.3, 0.5], dtype=tf.float32)

@tf.custom_gradient
def test_func_with_stop_grad(x):
  @tf.function
  def _inner_func():
    # Avoid exception during the forward pass
    return tf.stop_gradient(tf.random.shuffle(x))
    # return tf.random.shuffle(x)  # This will raise

  res = _inner_func()
  def grad(upstream):
    return upstream  # Arbitrarily defined custom gradient
  return res, grad

with tf.GradientTape() as g:
  g.watch(x)
  res = test_func_with_stop_grad(x)

g.gradient(res, x)

另請參閱tf.RegisterGradient，它為原始 TensorFlow 操作注冊梯度函數。另一方麵，tf.custom_gradient 允許對一係列操作的梯度計算進行細粒度控製。

請注意，如果修飾函數使用 Variable s，則封閉變量範圍必須使用 ResourceVariable s。