【tensorflow2.0】AutoGraph和tf.Module

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

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

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

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

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

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

前面我們介紹了Autograph的編碼規范和Autograph轉換成靜態圖的原理。

本篇我們介紹使用tf.Module來更好地構建Autograph。

一，Autograph和tf.Module概述

前面在介紹Autograph的編碼規范時提到構建Autograph時應該避免在@tf.function修飾的函數內部定義tf.Variable.

但是如果在函數外部定義tf.Variable的話，又會顯得這個函數有外部變量依賴，封裝不夠完美。

一種簡單的思路是定義一個類，並將相關的tf.Variable創建放在類的初始化方法中。而將函數的邏輯放在其他方法中。

這樣一頓猛如虎的操作之后，我們會覺得一切都如同人法地地法天天法道道法自然般的自然。

驚喜的是，TensorFlow提供了一個基類tf.Module，通過繼承它構建子類，我們不僅可以獲得以上的自然而然，而且可以非常方便地管理變量，還可以非常方便地管理它引用的其它Module，最重要的是，我們能夠利用tf.saved_model保存模型並實現跨平台部署使用。

實際上，tf.keras.models.Model,tf.keras.layers.Layer 都是繼承自tf.Module的，提供了方便的變量管理和所引用的子模塊管理的功能。

因此，利用tf.Module提供的封裝，再結合TensoFlow豐富的低階API，實際上我們能夠基於TensorFlow開發任意機器學習模型(而非僅僅是神經網絡模型)，並實現跨平台部署使用。

二，應用tf.Module封裝Autograph

定義一個簡單的function。

import tensorflow as tf 
x = tf.Variable(1.0,dtype=tf.float32)
 
# 在tf.function中用input_signature限定輸入張量的簽名類型：shape和dtype
@tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.float32)])    
def add_print(a):
    x.assign_add(a)
    tf.print(x)
    return(x)
add_print(tf.constant(3.0))
# add_print(tf.constant(3)) #輸入不符合張量簽名的參數將報錯

4

下面利用tf.Module的子類化將其封裝一下。

class DemoModule(tf.Module):
    def __init__(self,init_value = tf.constant(0.0),name=None):
        super(DemoModule, self).__init__(name=name)
        with self.name_scope:  #相當於with tf.name_scope("demo_module")
            self.x = tf.Variable(init_value,dtype = tf.float32,trainable=True)
 
 
    @tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.float32)])  
    def addprint(self,a):
        with self.name_scope:
            self.x.assign_add(a)
            tf.print(self.x)
            return(self.x)
 
# 執行
demo = DemoModule(init_value = tf.constant(1.0))
result = demo.addprint(tf.constant(5.0))

6

# 查看模塊中的全部變量和全部可訓練變量
print(demo.variables)
print(demo.trainable_variables)

(<tf.Variable 'demo_module/Variable:0' shape=() dtype=float32, numpy=6.0>,)
(<tf.Variable 'demo_module/Variable:0' shape=() dtype=float32, numpy=6.0>,)

# 查看模塊中的全部子模塊
demo.submodules

# 使用tf.saved_model 保存模型，並指定需要跨平台部署的方法
tf.saved_model.save(demo,"./data/demo/1",signatures = {"serving_default":demo.addprint})

WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
INFO:tensorflow:Assets written to: ./data/demo/1/assets

# 加載模型
demo2 = tf.saved_model.load("./data/demo/1")
demo2.addprint(tf.constant(5.0))

11

<tf.Tensor: shape=(), dtype=float32, numpy=11.0>

# 查看模型文件相關信息，紅框標出來的輸出信息在模型部署和跨平台使用時有可能會用到
!saved_model_cli show --dir ./data/demo/1 --all

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['__saved_model_init_op']:
  The given SavedModel SignatureDef contains the following input(s):
  The given SavedModel SignatureDef contains the following output(s):
    outputs['__saved_model_init_op'] tensor_info:
        dtype: DT_INVALID
        shape: unknown_rank
        name: NoOp
  Method name is: 

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['a'] tensor_info:
        dtype: DT_FLOAT
        shape: ()
        name: serving_default_a:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output_0'] tensor_info:
        dtype: DT_FLOAT
        shape: ()
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict
WARNING: Logging before flag parsing goes to stderr.
W0411 02:47:30.452981 139671888869248 deprecation.py:506] From /usr/local/lib/python2.7/dist-packages/tensorflow_core/python/ops/resource_variable_ops.py:1786: calling __init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.

Defined Functions:
  Function Name: 'addprint'
    Option #1
      Callable with:
        Argument #1
          a: TensorSpec(shape=(), dtype=tf.float32, name=u'a')

在tensorboard中查看計算圖，模塊會被添加模塊名demo_module,方便層次化呈現計算圖結構。

import datetime
 
# 創建日志
stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = './data/demomodule/%s' % stamp
writer = tf.summary.create_file_writer(logdir)
 
# 開啟autograph跟蹤
tf.summary.trace_on(graph=True, profiler=True) 
 
# 執行autograph
demo = DemoModule(init_value = tf.constant(0.0))
result = demo.addprint(tf.constant(5.0))
 
# 將計算圖信息寫入日志
with writer.as_default():
    tf.summary.trace_export(
        name="demomodule",
        step=0,
        profiler_outdir=logdir)
 
 
# 啟動 tensorboard在jupyter中的魔法命令
%reload_ext tensorboard
from tensorboard import notebook
notebook.list() 
notebook.start("--logdir ./data/demomodule/")

除了利用tf.Module的子類化實現封裝，我們也可以通過給tf.Module添加屬性的方法進行封裝。 

mymodule = tf.Module()
mymodule.x = tf.Variable(0.0)
 
@tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.float32)])  
def addprint(a):
    mymodule.x.assign_add(a)
    tf.print(mymodule.x)
    return (mymodule.x)
 
mymodule.addprint = addprint
mymodule.addprint(tf.constant(1.0)).numpy()

print(mymodule.variables)

# 使用tf.saved_model 保存模型
tf.saved_model.save(mymodule,"./data/mymodule",
    signatures = {"serving_default":mymodule.addprint})
 
# 加載模型
mymodule2 = tf.saved_model.load("./data/mymodule")
mymodule2.addprint(tf.constant(5.0))

1
(<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>,)
INFO:tensorflow:Assets written to: ./data/mymodule/assets
6

<tf.Tensor: shape=(), dtype=float32, numpy=6.0>

三，tf.Module和tf.keras.Model，tf.keras.layers.Layer

tf.keras中的模型和層都是繼承tf.Module實現的，也具有變量管理和子模塊管理功能。

import tensorflow as tf
from tensorflow.keras import models,layers,losses,metrics
print(issubclass(tf.keras.Model,tf.Module))
print(issubclass(tf.keras.layers.Layer,tf.Module))
print(issubclass(tf.keras.Model,tf.keras.layers.Layer))

True

True

True

tf.keras.backend.clear_session() 
 
model = models.Sequential()
 
model.add(layers.Dense(4,input_shape = (10,)))
model.add(layers.Dense(2))
model.add(layers.Dense(1))
model.summary()

model.variables

[<tf.Variable 'dense/kernel:0' shape=(10, 4) dtype=float32, numpy=
 array([[-0.24266458, -0.45152673, -0.5430875 , -0.35098866],
        [ 0.36108053, -0.32325   ,  0.3329792 ,  0.33279514],
        [ 0.2944306 , -0.5975202 , -0.06157887,  0.25049144],
        [ 0.5707406 ,  0.6214677 , -0.32870707, -0.12539297],
        [ 0.41170907, -0.5257766 ,  0.12482923, -0.11132008],
        [-0.41743976, -0.3998926 , -0.46740663,  0.6105366 ],
        [ 0.54347396,  0.5108323 ,  0.4747305 , -0.404514  ],
        [ 0.4390788 , -0.1988923 ,  0.40562296,  0.57931125],
        [-0.2694599 , -0.4149857 ,  0.07898462, -0.05845898],
        [-0.02557009, -0.440827  , -0.26627067, -0.0769726 ]],
       dtype=float32)>,
 <tf.Variable 'dense/bias:0' shape=(4,) dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>,
 <tf.Variable 'dense_1/kernel:0' shape=(4, 2) dtype=float32, numpy=
 array([[ 0.17386723,  0.9130187 ],
        [-0.88832307, -0.20379901],
        [ 0.9303725 , -0.4667368 ],
        [-0.8743646 , -0.31934786]], dtype=float32)>,
 <tf.Variable 'dense_1/bias:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>,
 <tf.Variable 'dense_2/kernel:0' shape=(2, 1) dtype=float32, numpy=
 array([[ 0.4870274 ],
        [-0.71679246]], dtype=float32)>,
 <tf.Variable 'dense_2/bias:0' shape=(1,) dtype=float32, numpy=array([0.], dtype=float32)>]

model.layers[0].trainable = False #凍結第0層的變量,使其不可訓練
model.trainable_variables

[<tf.Variable 'dense_1/kernel:0' shape=(4, 2) dtype=float32, numpy=
 array([[ 0.17386723,  0.9130187 ],
        [-0.88832307, -0.20379901],
        [ 0.9303725 , -0.4667368 ],
        [-0.8743646 , -0.31934786]], dtype=float32)>,
 <tf.Variable 'dense_1/bias:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>,
 <tf.Variable 'dense_2/kernel:0' shape=(2, 1) dtype=float32, numpy=
 array([[ 0.4870274 ],
        [-0.71679246]], dtype=float32)>,
 <tf.Variable 'dense_2/bias:0' shape=(1,) dtype=float32, numpy=array([0.], dtype=float32)>]

model.submodules

(<tensorflow.python.keras.engine.input_layer.InputLayer at 0x7fac6c6c2278>,
 <tensorflow.python.keras.layers.core.Dense at 0x7fac6e3e1908>,
 <tensorflow.python.keras.layers.core.Dense at 0x7fac6c6c2438>,
 <tensorflow.python.keras.layers.core.Dense at 0x7fac6c6c2470>)

model.layers

[<tensorflow.python.keras.layers.core.Dense at 0x7fac6e3e1908>,
 <tensorflow.python.keras.layers.core.Dense at 0x7fac6c6c2438>,
 <tensorflow.python.keras.layers.core.Dense at 0x7fac6c6c2470>]

print(model.name)
print(model.name_scope())

sequential
sequential

 

 

參考：

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

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