【tensorflow2.0】AutoGraph的機制原理

2020 年 4 月 11 日
筆記

有三種計算圖的構建方式：靜態計算圖，動態計算圖，以及Autograph。

TensorFlow 2.0主要使用的是動態計算圖和Autograph。

動態計算圖易於調試，編碼效率較高，但執行效率偏低。

靜態計算圖執行效率很高，但較難調試。

而Autograph機制可以將動態圖轉換成靜態計算圖，兼收執行效率和編碼效率之利。

當然Autograph機制能夠轉換的代碼並不是沒有任何約束的，有一些編碼規範需要遵循，否則可能會轉換失敗或者不符合預期。

我們會介紹Autograph的編碼規範和Autograph轉換成靜態圖的原理。

並介紹使用tf.Module來更好地構建Autograph。

上篇我們介紹了Autograph的編碼規範，本篇我們介紹Autograph的機制原理。

一，Autograph的機制原理

當我們使用@tf.function裝飾一個函數的時候，後面到底發生了什麼呢？

例如我們寫下如下代碼。

import tensorflow as tf  import numpy as np    @tf.function(autograph=True)  def myadd(a,b):      for i in tf.range(3):          tf.print(i)      c = a+b      print("tracing")      return c

後面什麼都沒有發生。僅僅是在Python堆棧中記錄了這樣一個函數的簽名。

當我們第一次調用這個被@tf.function裝飾的函數時，後面到底發生了什麼？

例如我們寫下如下代碼。

myadd(tf.constant("hello"),tf.constant("world"))

tracing  0  1  2

<tf.Tensor: shape=(), dtype=string, numpy=b'helloworld'>

發生了2件事情，

第一件事情是創建計算圖。

即創建一個靜態計算圖，跟蹤執行一遍函數體中的Python代碼，確定各個變量的Tensor類型，並根據執行順序將算子添加到計算圖中。在這個過程中，如果開啟了autograph=True(默認開啟),會將Python控制流轉換成TensorFlow圖內控制流。主要是將if語句轉換成 tf.cond算子表達，將while和for循環語句轉換成tf.while_loop算子表達，並在必要的時候添加 tf.control_dependencies指定執行順序依賴關係。

相當於在 tensorflow1.0執行了類似下面的語句：

g = tf.Graph()  with g.as_default():      a = tf.placeholder(shape=[],dtype=tf.string)      b = tf.placeholder(shape=[],dtype=tf.string)      cond = lambda i: i<tf.constant(3)      def body(i):          tf.print(i)          return(i+1)      loop = tf.while_loop(cond,body,loop_vars=[0])      loop      with tf.control_dependencies(loop):          c = tf.strings.join([a,b])      print("tracing")

第二件事情是執行計算圖。

相當於在 tensorflow1.0中執行了下面的語句：

with tf.Session(graph=g) as sess:      sess.run(c,feed_dict={a:tf.constant("hello"),b:tf.constant("world")})

因此我們先看到的是第一個步驟的結果：即Python調用標準輸出流打印”tracing”語句。

然後看到第二個步驟的結果：TensorFlow調用標準輸出流打印1,2,3。

當我們再次用相同的輸入參數類型調用這個被@tf.function裝飾的函數時，後面到底發生了什麼？

例如我們寫下如下代碼。

myadd(tf.constant("hello"),tf.constant("world"))

0  1  2

<tf.Tensor: shape=(), dtype=string, numpy=b'helloworld'>

只會發生一件事情，那就是上面步驟的第二步，執行計算圖。

所以這一次我們沒有看到打印”tracing”的結果。

當我們再次用不同的的輸入參數類型調用這個被@tf.function裝飾的函數時，後面到底發生了什麼？

例如我們寫下如下代碼。

myadd(tf.constant(1),tf.constant(2))

tracing  0  1  2

<tf.Tensor: shape=(), dtype=int32, numpy=3>

由於輸入參數的類型已經發生變化，已經創建的計算圖不能夠再次使用。

需要重新做2件事情：創建新的計算圖、執行計算圖。

所以我們又會先看到的是第一個步驟的結果：即Python調用標準輸出流打印”tracing”語句。

然後再看到第二個步驟的結果：TensorFlow調用標準輸出流打印1,2,3。

需要注意的是，如果調用被@tf.function裝飾的函數時輸入的參數不是Tensor類型，則每次都會重新創建計算圖。

例如我們寫下如下代碼。兩次都會重新創建計算圖。因此，一般建議調用@tf.function時應傳入Tensor類型。

myadd("hello","world")  myadd("good","morning")

tracing

二，重新理解Autograph的編碼規範

了解了以上Autograph的機制原理，我們也就能夠理解Autograph編碼規範的3條建議了。

1，被@tf.function修飾的函數應盡量使用TensorFlow中的函數而不是Python中的其他函數。例如使用tf.print而不是print.

解釋：Python中的函數僅僅會在跟蹤執行函數以創建靜態圖的階段使用，普通Python函數是無法嵌入到靜態計算圖中的，所以在計算圖構建好之後再次調用的時候，這些Python函數並沒有被計算，而TensorFlow中的函數則可以嵌入到計算圖中。使用普通的Python函數會導致被@tf.function修飾前【eager執行】和被@tf.function修飾後【靜態圖執行】的輸出不一致。

2，避免在@tf.function修飾的函數內部定義tf.Variable.

解釋：如果函數內部定義了tf.Variable,那麼在【eager執行】時，這種創建tf.Variable的行為在每次函數調用時候都會發生。但是在【靜態圖執行】時，這種創建tf.Variable的行為只會發生在第一步跟蹤Python代碼邏輯創建計算圖時，這會導致被@tf.function修飾前【eager執行】和被@tf.function修飾後【靜態圖執行】的輸出不一致。實際上，TensorFlow在這種情況下一般會報錯。

3，被@tf.function修飾的函數不可修改該函數外部的Python列表或字典等數據結構變量。

解釋：靜態計算圖是被編譯成C++代碼在TensorFlow內核中執行的。Python中的列表和字典等數據結構變量是無法嵌入到計算圖中，它們僅僅能夠在創建計算圖時被讀取，在執行計算圖時是無法修改Python中的列表或字典這樣的數據結構變量的。

參考：

開源電子書地址：https://lyhue1991.github.io/eat_tensorflow2_in_30_days/

GitHub 項目地址：https://github.com/lyhue1991/eat_tensorflow2_in_30_days