給讀者：這是鐵人賽的搶先版，往後陸續還有更新。

在上一篇文章中我們提到了如何在 eager mode 2.0 處理tf.Variable 的眾多問題。今天我們將會介紹如何利用tf.Variable 和 tf.GradientTape 在 keras 的模型中客製化梯度。現在就來看一個例子：
首先，讓我們建立一個 keras 模型，這個模型有兩個 2D convolution 層，一個 2D GlobalAveragePooling，最後則是一個全連階層。

# Build the model
mnist_model = tf.keras.Sequential([
  tf.keras.layers.Conv2D(16,[3,3], activation='relu',
                         input_shape=(None, None, 1)),
  tf.keras.layers.Conv2D(16,[3,3], activation='relu'),
  tf.keras.layers.GlobalAveragePooling2D(),
  tf.keras.layers.Dense(10)
])

雖然 keras 的模型可以使用函式庫所提供的 training loop，但為能夠客製化訓練，於是我們自行寫一個 training loop，如下：

def train_step(images, labels):
  with tf.GradientTape() as tape: # 開始記錄正向傳播計
    logits = mnist_model(images, training=True)
    
    #  tf.debuggin 可以用在 eager mode 和 graph mode
    # 加入 asserts 用以檢查輸出變數的維度
    tf.debugging.assert_equal(logits.shape, (32, 10))
    
    loss_value = loss_object(labels, logits)

  loss_history.append(loss_value.numpy().mean())
  # 計算 gradient 反向傳播
  grads = tape.gradient(loss_value, mnist_model.trainable_variables)
  optimizer.apply_gra
  dients(zip(grads, mnist_model.trainable_variables))

從上面的程式碼可以看到，tf.GradientTape()可以用在 training loop 裡，紀錄並建構正向傳播的計算圖，在完成“記錄”後，tf.GradientTape() 的 tape 物件，宛如PyTorch 的損失張量呼叫 backward()方法， tape 則呼叫 gradient()的方法，並傳入損失張量和模型可訓練的參數。
一旦計算出了梯度後，立即呼叫optimizer.apply_gradients()方法，傳入一個 list of tuple，每一個 tuple 的第二個則是參數變數，而第一個變數為針對該參數所計算出的梯度。這個函式在 PyTorch 則是 optimizer.step()。
最後，這個過程無需損失張量去清空 buffer，也就是 loss.zero()的方法，對於不擅長管理記憶體的人，這個設計非常有幫助。

tf.custom_gradient

接下來我們要介紹 tf.custom_gradient()這個函式。這個函式其實可以當作 decorator 來使用，也就是用來包覆可以客製化梯度的函式，並在函式執行的前後執行某些檢查或後處理等程式邏輯。
簡單地說，若要手動呼叫這個函式，就需要傳入一個處理梯度的函式，f，而呼叫方法則是tf.custom_gradient(f)。但大部分的時候，則是用 decorator 的方式，在 f 的定義上頭加上 @tf.custom_gradient。
客製化梯度給予使用者可以對梯度計算做細度的控制，如對梯度做 clip，以防止梯度爆炸的問題。下面的程式範例就是使用梯度客製化函式來解決數值不穩定的問題。在程式中，我們先定義一個函式，log1pexp。在這個函式裡，主要是對公式做計算：。然而，這個函式的輸入為較大的數值時，會產生接近無限的 exp(x)輸出，而因此為數值不穩定。我們可以先看這個數值不穩定的版本：

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

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

但我們可以藉著 tf.custom_gradient 來提供客製化的梯度計算，避免掉這個問題。

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