iT邦幫忙

第 11 屆 iThome 鐵人賽

DAY 17
1
AI & Data

深度學習裡的冰與火之歌 : Tensorflow vs PyTorch系列 第 17

Day 17: Tensorflow 2.0 再造訪 tf.GradientTape

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

上一篇文章中我們提到了如何在 eager mode 2.0 處理tf.Variable 的眾多問題。今天我們將會介紹如何利用tf.Variabletf.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。在這個函式裡,主要是對公式做計算:https://chart.googleapis.com/chart?cht=tx&chl=%5Clog(1%20%2B%20%5Cexp(x))。然而,這個函式的輸入為較大的數值時,會產生接近無限的 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

上一篇
Day 16: Tensorflow 2.0 再造訪 eager-execution
下一篇
Day 18: Tensorflow 2.0 再造訪 tf.function
系列文
深度學習裡的冰與火之歌 : Tensorflow vs PyTorch31

尚未有邦友留言

立即登入留言