Day 19: Tensorflow 2.0 再造訪 AutoGraph

11th鐵人賽

renewang

2019-10-04 21:08:26

2274 瀏覽

分享至

在上一篇文章中稍微提到了 tf.function 在 Tensorflow 2.0 的角色，這次我們仍舊會著重在 tf.function，但也會介紹 tf.function 的底層實作細節，那也就是 AutoGraph。Tensorflow 另外有一個 tf.autograph 模組，主要是提供使用者介面，可以 access AutoGraph 的一些特徵，如上一篇文章的例子。

雖然說，AutoGraph 的內容都實作在 tf.function 那麼倒底這個 AutoGraph 需要實作些什麼東西呢？

為了偵測需要修改的地方， AutoGraph 必須分析原始碼（AutoGraph analyzes code to detect modifications）：在上一篇文章提到，AutoGraph 會將 python 語法轉譯成 tensorflow 的語法。事實上，AutoGraph 必須要先對 Python 原始碼做靜態原始碼分析（static analysis），才能決定該做如何的修改才能保持所建立的靜態圖是高效能的，修改的重點包括在 python 語法內，找出 loop_vars

y = 0
while x > 0:
  x = x - 1
print(y)
# 對應的 tensorflow 語法：
# x = tf.while_loop(
# cond=lambda x: x > 0,
# body=lambda x: x-=1,
# loop_vars=(x,) 
# # y 不是 loop_vars

在 Tensorflow 的 control flow 處理 Python collections（Python collections in TensorFlow control flow）：Python collections 包括了 list, tuple 和 dict ，因為長度有變動性所以和 Tensorflow 的 control flow 不相容。如果要使用 Python collections，甚至 python class。替代方案，包括了使用 tf.TensorArray(tf.int32, size=0, dynamic_size=True)來替換 list，
在 Tensorflow control flow 保持陣列維度與資料型態一致（Shape and dtype consistency in TensorFlow control flow）： Tensorflow 只有讚少數情況支持 dynamic typing，然而 Python 是 strict （檢查 type consistent）但 dynamic typing （由 context 決定 typing）。為了能符合 Tensorflow control flow，ＡutoGraph 的實作必須能偵測不一致之處。
處理與 Tensorflow 不相容的未定義以及 None 物件（Undefined and None values in TensorFlow）：未定義以及 None 物件是允許出現在 python control flow 因此在 ** eager mode** 允許，但這些物件在 Tensorflow 的 control flow 不被允許，AutoGraph 需要有能力處理這個意外。
AutoGraph 只能作用在 function 它的 source code 可以被取得（Access to source code）：包括 labmda function

control dependencies again

在[上一篇文章](https://ith1l elp.ithome.com.tw/articles/10224165)中提到

a = tf.Variable(1.0)
b = tf.Variable(1.0)

@tf.function
def f():
  # a.assign 和 b.assign 的執行順序，使用者無法控制
  a.assign(2.0) 
  b.assign(3.0)
  return a + b
  
print(f())

Automatic Control Dependencies

有 side-effect 的運算元：tf.assign 和 tf.print 等運算元。這些運算元被包覆在 tf.function 的時候，在執行前會先自動設立執行的相依關係，最終自動決定執行的順序。
在底下的例子中，第一個 assign op 為第二個 assign op 的依賴者，也就是第二個 assign op 必須要依賴第一個 assign op 的結果來執行，所以相依關係必須要強制包括進計算圖。

v = tf.Variable(5)

@tf.function
def find_next_odd():
  v.assign(v + 1) # v += 1, in-place addition
  if v % 2 == 0: 
    v.assign(v + 1) # 依賴 v 目前的值，故不能先執行

find_next_odd()
v

讀者若有興趣知道 automatic control dependencies 是如何執行，可以閱讀這段原始碼

keras 和 AutoGraph

什麼是 dynamic model in keras

在Tensorflow 2.0 再造訪 eager-execution 有提到若在__init__方法的 dynamic 引數設為 True，則會建立 dynamic model，以指令式編程（imperative programming）的方式來執行模型。不過 tf.function是屬於事先編譯的程式模型所以在這裡必須將 dynamic 引數設為 False。
現在，就讓我們看一下 non-dynamic 的 tf.keras.Model 如何和 tf.function結合：

class CustomModel(tf.keras.models.Model):

  @tf.function
  def call(self, input_data):
    if tf.reduce_mean(input_data) > 0:
      return input_data
    else:
      return input_data // 2

model = CustomModel()
model(tf.constant([-2, -4]))
#=> <tf.Tensor: id=710, shape=(2,), dtype=int32, numpy=array([-1, -2], dtype=int32)>

An in-graph training loop

下面的程式碼是使用 tf.function 和 tf.GradientTape 來建立 training loop。
在進入 training loop 之前，須用 Sequential API 建立好一個簡單的 keras 模型，而 optimizer 則是 tf.keras.optimizers.Adam() 建立。
資料的部分則是用 tf.keras.datasets.mnist.load_data() 來 load MNIST 手寫數字影像集到記憶體內，並用 tf.data.Dataset.from_tensor_slices 來轉成 tf.data.Dataset，batch size 為 100。
從這個範例我們可以看到整個 trainingn process 包括了 tf.data 將資料批次載入，計算梯度，更新參數，計算驗證集正確率，以及重複整個過程直到收斂為止，都可以在 tf.function的視野下，根據計算圖執行。

compute_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

compute_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()

def train_one_step(model, optimizer, x, y):
  with tf.GradientTape() as tape:
    logits = model(x)
    loss = compute_loss(y, logits)

  grads = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(grads, model.trainable_variables))

  compute_accuracy(y, logits)
  return loss


@tf.function
def train(model, optimizer):
  train_ds = mnist_dataset()
  step = 0
  loss = 0.0
  accuracy = 0.0
  for x, y in train_ds:
    step += 1
    loss = train_one_step(model, optimizer, x, y)
    if step % 10 == 0:
      tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())
  return step, loss, accuracy

step, loss, accuracy = train(model, optimizer)
print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())

上述的程式碼，大家可以試試看結果，這裡就不再把輸出列出。