梯度與自動微分簡介

在 TensorFlow.org 上查看

在 Google Colab 中執行

在 GitHub 上查看原始碼

下載筆記本

自動微分和梯度

自動微分適用於實作機器學習演算法，例如用於訓練神經網路的反向傳播。

在本指南中，您將探索使用 TensorFlow 計算梯度的方法，尤其是在立即執行中。

設定

import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf

計算梯度

若要自動微分，TensorFlow 必須記住前向傳遞期間以何種順序執行哪些運算。然後，在反向傳遞期間，TensorFlow 會反向遍歷此運算清單，以計算梯度。

梯度帶

TensorFlow 提供 tf.GradientTape API 用於自動微分；也就是計算運算相對於某些輸入 (通常是 tf.Variable) 的梯度。TensorFlow 會將在 tf.GradientTape 環境定義中執行的相關運算「記錄」到「磁帶」上。然後，TensorFlow 會使用該磁帶，透過反向模式微分來計算「已記錄」運算的梯度。

以下是一個簡單範例

x = tf.Variable(3.0)

with tf.GradientTape() as tape:
  y = x**2

記錄一些運算後，請使用 GradientTape.gradient(target, sources) 來計算目標 (通常是損失) 相對於來源 (通常是模型的變數) 的梯度

# dy = 2x * dx
dy_dx = tape.gradient(y, x)
dy_dx.numpy()

以上範例使用純量，但 tf.GradientTape 同樣適用於任何張量

w = tf.Variable(tf.random.normal((3, 2)), name='w')
b = tf.Variable(tf.zeros(2, dtype=tf.float32), name='b')
x = [[1., 2., 3.]]

with tf.GradientTape(persistent=True) as tape:
  y = x @ w + b
  loss = tf.reduce_mean(y**2)

若要取得 loss 相對於這兩個變數的梯度，您可以將這兩個變數都做為來源傳遞至 gradient 方法。磁帶對於來源的傳遞方式具有彈性，並接受清單或字典的任何巢狀組合，並以相同的結構傳回梯度 (請參閱 tf.nest)。

[dl_dw, dl_db] = tape.gradient(loss, [w, b])

相對於每個來源的梯度都具有來源的形狀

print(w.shape)
print(dl_dw.shape)

以下再次進行梯度計算，這次傳遞變數字典

my_vars = {
    'w': w,
    'b': b
}

grad = tape.gradient(loss, my_vars)
grad['b']

模型相對於模型的梯度

常見的做法是將 tf.Variables 收集到 tf.Module 或其子類別 (layers.Layer、keras.Model) 之中，以進行檢查點和匯出。

在大多數情況下，您會想要計算模型的可訓練變數的梯度。由於 tf.Module 的所有子類別都會在其 Module.trainable_variables 屬性中彙總其變數，因此您可以用幾行程式碼計算這些梯度

layer = tf.keras.layers.Dense(2, activation='relu')
x = tf.constant([[1., 2., 3.]])

with tf.GradientTape() as tape:
  # Forward pass
  y = layer(x)
  loss = tf.reduce_mean(y**2)

# Calculate gradients with respect to every trainable variable
grad = tape.gradient(loss, layer.trainable_variables)

for var, g in zip(layer.trainable_variables, grad):
  print(f'{var.name}, shape: {g.shape}')

控制磁帶監看的內容

預設行為是存取可訓練的 tf.Variable 後，記錄所有運算。原因如下：

• 磁帶需要知道在前向傳遞中要記錄哪些運算，才能在反向傳遞中計算梯度。
• 磁帶會保留對中繼輸出的參照，因此您不會想要記錄不必要的運算。
• 最常見的用途是計算損失相對於模型所有可訓練變數的梯度。

例如，以下範例無法計算梯度，因為 tf.Tensor 預設為「未受監看」，且 tf.Variable 不可訓練

# A trainable variable
x0 = tf.Variable(3.0, name='x0')
# Not trainable
x1 = tf.Variable(3.0, name='x1', trainable=False)
# Not a Variable: A variable + tensor returns a tensor.
x2 = tf.Variable(2.0, name='x2') + 1.0
# Not a variable
x3 = tf.constant(3.0, name='x3')

with tf.GradientTape() as tape:
  y = (x0**2) + (x1**2) + (x2**2)

grad = tape.gradient(y, [x0, x1, x2, x3])

for g in grad:
  print(g)

您可以使用 GradientTape.watched_variables 方法列出磁帶正在監看的變數

[var.name for var in tape.watched_variables()]

tf.GradientTape 提供掛鉤，讓使用者可以控制要監看或不監看的內容。

若要記錄相對於 tf.Tensor 的梯度，您需要呼叫 GradientTape.watch(x)

x = tf.constant(3.0)
with tf.GradientTape() as tape:
  tape.watch(x)
  y = x**2

# dy = 2x * dx
dy_dx = tape.gradient(y, x)
print(dy_dx.numpy())

反之，若要停用監看所有 tf.Variables 的預設行為，請在建立梯度帶時設定 watch_accessed_variables=False。此計算使用兩個變數，但只連線其中一個變數的梯度

x0 = tf.Variable(0.0)
x1 = tf.Variable(10.0)

with tf.GradientTape(watch_accessed_variables=False) as tape:
  tape.watch(x1)
  y0 = tf.math.sin(x0)
  y1 = tf.nn.softplus(x1)
  y = y0 + y1
  ys = tf.reduce_sum(y)

由於未在 x0 上呼叫 GradientTape.watch，因此未計算相對於它的梯度

# dys/dx1 = exp(x1) / (1 + exp(x1)) = sigmoid(x1)
grad = tape.gradient(ys, {'x0': x0, 'x1': x1})

print('dy/dx0:', grad['x0'])
print('dy/dx1:', grad['x1'].numpy())

中繼結果

您也可以要求輸出相對於 tf.GradientTape 環境定義內計算的中繼值的梯度。

x = tf.constant(3.0)

with tf.GradientTape() as tape:
  tape.watch(x)
  y = x * x
  z = y * y

# Use the tape to compute the gradient of z with respect to the
# intermediate value y.
# dz_dy = 2 * y and y = x ** 2 = 9
print(tape.gradient(z, y).numpy())

預設情況下，GradientTape 持有的資源會在呼叫 GradientTape.gradient 方法後立即釋出。若要針對相同的運算計算多個梯度，請使用 persistent=True 建立梯度帶。這允許多次呼叫 gradient 方法，因為資源會在垃圾收集磁帶物件時釋出。例如：

x = tf.constant([1, 3.0])
with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  y = x * x
  z = y * y

print(tape.gradient(z, x).numpy())  # [4.0, 108.0] (4 * x**3 at x = [1.0, 3.0])
print(tape.gradient(y, x).numpy())  # [2.0, 6.0] (2 * x at x = [1.0, 3.0])

del tape   # Drop the reference to the tape

效能注意事項

• 在梯度帶環境定義內執行運算會有輕微的額外負擔。對於大多數立即執行而言，這不會是明顯的成本，但您仍應僅在需要的位置周圍使用磁帶環境定義。

• 梯度帶會使用記憶體來儲存中繼結果，包括輸入和輸出，以便在反向傳遞期間使用。

  為了提高效率，某些運算 (例如 ReLU) 不需要保留其中繼結果，而且這些結果會在前向傳遞期間修剪。不過，如果您在磁帶上使用 persistent=True，則不會捨棄任何項目，而且您的尖峰記憶體用量會更高。

非純量目標的梯度

梯度基本上是純量上的運算。

x = tf.Variable(2.0)
with tf.GradientTape(persistent=True) as tape:
  y0 = x**2
  y1 = 1 / x

print(tape.gradient(y0, x).numpy())
print(tape.gradient(y1, x).numpy())

因此，如果您要求多個目標的梯度，則每個來源的結果為

• 目標總和的梯度，或等效地說
• 每個目標的梯度總和。

x = tf.Variable(2.0)
with tf.GradientTape() as tape:
  y0 = x**2
  y1 = 1 / x

print(tape.gradient({'y0': y0, 'y1': y1}, x).numpy())

同樣地，如果目標不是純量，則會計算總和的梯度

x = tf.Variable(2.)

with tf.GradientTape() as tape:
  y = x * [3., 4.]

print(tape.gradient(y, x).numpy())

這讓您可以輕鬆取得損失集合總和的梯度，或元素方式損失計算總和的梯度。

如果您需要每個項目的個別梯度，請參閱雅可比行列式。

在某些情況下，您可以略過雅可比行列式。對於元素方式計算，總和的梯度會提供每個元素相對於其輸入元素的導數，因為每個元素都是獨立的

x = tf.linspace(-10.0, 10.0, 200+1)

with tf.GradientTape() as tape:
  tape.watch(x)
  y = tf.nn.sigmoid(x)

dy_dx = tape.gradient(y, x)

plt.plot(x, y, label='y')
plt.plot(x, dy_dx, label='dy/dx')
plt.legend()
_ = plt.xlabel('x')

控制流程

由於梯度帶會在運算執行時記錄運算，因此 Python 控制流程會自然地處理 (例如，if 和 while 陳述式)。

此處在 if 的每個分支上使用不同的變數。梯度只會連線至已使用的變數

x = tf.constant(1.0)

v0 = tf.Variable(2.0)
v1 = tf.Variable(2.0)

with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  if x > 0.0:
    result = v0
  else:
    result = v1**2 

dv0, dv1 = tape.gradient(result, [v0, v1])

print(dv0)
print(dv1)

只要記住控制陳述式本身不可微分，因此對以梯度為基礎的最佳化工具而言是看不見的。

在以上範例中，根據 x 的值，磁帶會記錄 result = v0 或 result = v1**2。x 相對於 x 的梯度一律為 None。

dx = tape.gradient(result, x)

print(dx)

gradient 傳回 None 的情況

當目標未連線至來源時，gradient 會傳回 None。

x = tf.Variable(2.)
y = tf.Variable(3.)

with tf.GradientTape() as tape:
  z = y * y
print(tape.gradient(z, x))

此處 z 明顯未連線至 x，但梯度可能會以幾種較不明顯的方式中斷連線。

1. 以張量取代變數

在「控制磁帶監看的內容」一節中，您已看到磁帶會自動監看 tf.Variable，但不會監看 tf.Tensor。

一個常見的錯誤是不小心以 tf.Tensor 取代 tf.Variable，而不是使用 Variable.assign 更新 tf.Variable。以下範例：

x = tf.Variable(2.0)

for epoch in range(2):
  with tf.GradientTape() as tape:
    y = x+1

  print(type(x).__name__, ":", tape.gradient(y, x))
  x = x + 1   # This should be `x.assign_add(1)`

2. 在 TensorFlow 外部執行計算

如果計算結束 TensorFlow，磁帶就無法記錄梯度路徑。例如：

x = tf.Variable([[1.0, 2.0],
                 [3.0, 4.0]], dtype=tf.float32)

with tf.GradientTape() as tape:
  x2 = x**2

  # This step is calculated with NumPy
  y = np.mean(x2, axis=0)

  # Like most ops, reduce_mean will cast the NumPy array to a constant tensor
  # using `tf.convert_to_tensor`.
  y = tf.reduce_mean(y, axis=0)

print(tape.gradient(y, x))

3. 透過整數或字串取得梯度

整數和字串不可微分。如果計算路徑使用這些資料類型，則不會有梯度。

沒有人會預期字串可微分，但如果您未指定 dtype，很容易不小心建立 int 常數或變數。

x = tf.constant(10)

with tf.GradientTape() as g:
  g.watch(x)
  y = x * x

print(g.gradient(y, x))

TensorFlow 不會自動在類型之間轉換，因此，實際上，您通常會收到類型錯誤，而不是遺失梯度。

4. 透過具狀態物件取得梯度

狀態會停止梯度。當您從具狀態物件讀取時，磁帶只能觀察目前狀態，而無法觀察導致該狀態的歷史記錄。

tf.Tensor 是不可變的。您無法在建立張量後變更張量。它有值，但沒有狀態。到目前為止討論的所有運算也都是無狀態的：tf.matmul 的輸出只取決於其輸入。

tf.Variable 具有內部狀態，即其值。當您使用變數時，會讀取狀態。計算相對於變數的梯度是正常的，但變數的狀態會阻擋梯度計算進一步回溯。例如：

x0 = tf.Variable(3.0)
x1 = tf.Variable(0.0)

with tf.GradientTape() as tape:
  # Update x1 = x1 + x0.
  x1.assign_add(x0)
  # The tape starts recording from x1.
  y = x1**2   # y = (x1 + x0)**2

# This doesn't work.
print(tape.gradient(y, x0))   #dy/dx0 = 2*(x1 + x0)

同樣地，tf.data.Dataset 迭代器和 tf.queue 具有狀態，而且會停止通過它們的張量的所有梯度。

未註冊梯度

某些 tf.Operation 會註冊為不可微分，並傳回 None。其他則未註冊梯度。

tf.raw_ops 頁面會顯示哪些低階運算已註冊梯度。

如果您嘗試透過未註冊梯度的浮點運算取得梯度，磁帶會擲回錯誤，而不是靜默傳回 None。這樣您就會知道發生錯誤。

例如，tf.image.adjust_contrast 函式會包裝 raw_ops.AdjustContrastv2，這可能會有梯度，但梯度尚未實作

image = tf.Variable([[[0.5, 0.0, 0.0]]])
delta = tf.Variable(0.1)

with tf.GradientTape() as tape:
  new_image = tf.image.adjust_contrast(image, delta)

try:
  print(tape.gradient(new_image, [image, delta]))
  assert False   # This should not happen.
except LookupError as e:
  print(f'{type(e).__name__}: {e}')

如果您需要透過此運算進行微分，則需要實作梯度並註冊 (使用 tf.RegisterGradient)，或使用其他運算重新實作函式。

零值而非 None

在某些情況下，對於未連線的梯度，取得 0 而非 None 會很方便。您可以使用 unconnected_gradients 引數決定在有未連線的梯度時要傳回的內容

x = tf.Variable([2., 2.])
y = tf.Variable(3.)

with tf.GradientTape() as tape:
  z = y**2
print(tape.gradient(z, x, unconnected_gradients=tf.UnconnectedGradients.ZERO))