4-5,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。

  1. import tensorflow as tf 
  2. x = tf.Variable(1.0,dtype=tf.float32)
  4. #在tf.function中用input_signature限定输入张量的签名类型：shape和dtype
  5. @tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.float32)])    
  6. def add_print(a):
  7.     x.assign_add(a)
  8.     tf.print(x)
  9.     return(x)
  1. add_print(tf.constant(3.0))
  2. #add_print(tf.constant(3)) #输入不符合张量签名的参数将报错
  1. 4

下面利用tf.Module的子类化将其封装一下。

  1. class DemoModule(tf.Module):
  2.     def __init__(self,init_value = tf.constant(0.0),name=None):
  3.         super(DemoModule, self).__init__(name=name)
  4.         with self.name_scope:  #相当于with tf.name_scope("demo_module")
  5.             self.x = tf.Variable(init_value,dtype = tf.float32,trainable=True)
  8.     @tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.float32)])  
  9.     def addprint(self,a):
  10.         with self.name_scope:
  11.             self.x.assign_add(a)
  12.             tf.print(self.x)
  13.             return(self.x)
  1. #执行
  2. demo = DemoModule(init_value = tf.constant(1.0))
  3. result = demo.addprint(tf.constant(5.0))
  1. 6
  1. #查看模块中的全部变量和全部可训练变量
  2. print(demo.variables)
  3. print(demo.trainable_variables)
  1. (<tf.Variable 'demo_module/Variable:0' shape=() dtype=float32, numpy=6.0>,)
  2. (<tf.Variable 'demo_module/Variable:0' shape=() dtype=float32, numpy=6.0>,)
  1. #查看模块中的全部子模块
  2. demo.submodules
  1. #使用tf.saved_model 保存模型，并指定需要跨平台部署的方法
  2. tf.saved_model.save(demo,"./data/demo/1",signatures = {"serving_default":demo.addprint})
  1. #加载模型
  2. demo2 = tf.saved_model.load("./data/demo/1")
  3. demo2.addprint(tf.constant(5.0))
  1. 11
  1. # 查看模型文件相关信息，红框标出来的输出信息在模型部署和跨平台使用时有可能会用到
  2. !saved_model_cli show --dir ./data/demo/1 --all

4-5,AutoGraph和tf.Module - 图1

在tensorboard中查看计算图，模块会被添加模块名demo_module,方便层次化呈现计算图结构。

  1. import datetime
  3. # 创建日志
  4. stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
  5. logdir = './data/demomodule/%s' % stamp
  6. writer = tf.summary.create_file_writer(logdir)
  8. #开启autograph跟踪
  9. tf.summary.trace_on(graph=True, profiler=True) 
  11. #执行autograph
  12. demo = DemoModule(init_value = tf.constant(0.0))
  13. result = demo.addprint(tf.constant(5.0))
  15. #将计算图信息写入日志
  16. with writer.as_default():
  17.     tf.summary.trace_export(
  18.         name="demomodule",
  19.         step=0,
  20.         profiler_outdir=logdir)
  1. #启动 tensorboard在jupyter中的魔法命令
  2. %reload_ext tensorboard
  1. from tensorboard import notebook
  2. notebook.list()
  1. notebook.start("--logdir ./data/demomodule/")

4-5,AutoGraph和tf.Module - 图2

除了利用tf.Module的子类化实现封装，我们也可以通过给tf.Module添加属性的方法进行封装。

  1. mymodule = tf.Module()
  2. mymodule.x = tf.Variable(0.0)
  4. @tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.float32)])  
  5. def addprint(a):
  6.     mymodule.x.assign_add(a)
  7.     tf.print(mymodule.x)
  8.     return (mymodule.x)
  10. mymodule.addprint = addprint
  1. mymodule.addprint(tf.constant(1.0)).numpy()
  1. 1.0
  1. print(mymodule.variables)
  1. (<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.0>,)
  1. #使用tf.saved_model 保存模型
  2. tf.saved_model.save(mymodule,"./data/mymodule",
  3.     signatures = {"serving_default":mymodule.addprint})
  5. #加载模型
  6. mymodule2 = tf.saved_model.load("./data/mymodule")
  7. mymodule2.addprint(tf.constant(5.0))
  1. INFO:tensorflow:Assets written to: ./data/mymodule/assets
  2. 5

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

tf.keras中的模型和层都是继承tf.Module实现的，也具有变量管理和子模块管理功能。

  1. import tensorflow as tf
  2. from tensorflow.keras import models,layers,losses,metrics
  1. print(issubclass(tf.keras.Model,tf.Module))
  2. print(issubclass(tf.keras.layers.Layer,tf.Module))
  3. print(issubclass(tf.keras.Model,tf.keras.layers.Layer))
  1. True
  2. True
  3. True
  1. tf.keras.backend.clear_session() 
  3. model = models.Sequential()
  5. model.add(layers.Dense(4,input_shape = (10,)))
  6. model.add(layers.Dense(2))
  7. model.add(layers.Dense(1))
  8. model.summary()
  1. Model: "sequential"
  2. _________________________________________________________________
  3. Layer (type)                 Output Shape              Param #   
  4. =================================================================
  5. dense (Dense)                (None, 4)                 44        
  6. _________________________________________________________________
  7. dense_1 (Dense)              (None, 2)                 10        
  8. _________________________________________________________________
  9. dense_2 (Dense)              (None, 1)                 3         
  10. =================================================================
  11. Total params: 57
  12. Trainable params: 57
  13. Non-trainable params: 0
  14. _________________________________________________________________
  1. model.variables
  1. [<tf.Variable 'dense/kernel:0' shape=(10, 4) dtype=float32, numpy=
  2.  array([[-0.06741005,  0.45534766,  0.5190817 , -0.01806331],
  3.         [-0.14258742, -0.49711505,  0.26030976,  0.18607801],
  4.         [-0.62806034,  0.5327399 ,  0.42206633,  0.29201728],
  5.         [-0.16602087, -0.18901917,  0.55159235, -0.01091868],
  6.         [ 0.04533798,  0.326845  , -0.582667  ,  0.19431782],
  7.         [ 0.6494713 , -0.16174704,  0.4062966 ,  0.48760796],
  8.         [ 0.58400524, -0.6280886 , -0.11265379, -0.6438277 ],
  9.         [ 0.26642334,  0.49275804,  0.20793378, -0.43889117],
  10.         [ 0.4092741 ,  0.09871006, -0.2073121 ,  0.26047975],
  11.         [ 0.43910992,  0.00199282, -0.07711256, -0.27966842]],
  12.        dtype=float32)>,
  13.  <tf.Variable 'dense/bias:0' shape=(4,) dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>,
  14.  <tf.Variable 'dense_1/kernel:0' shape=(4, 2) dtype=float32, numpy=
  15.  array([[ 0.5022683 , -0.0507431 ],
  16.         [-0.61540484,  0.9369011 ],
  17.         [-0.14412141, -0.54607415],
  18.         [ 0.2027781 , -0.4651153 ]], dtype=float32)>,
  19.  <tf.Variable 'dense_1/bias:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>,
  20.  <tf.Variable 'dense_2/kernel:0' shape=(2, 1) dtype=float32, numpy=
  21.  array([[-0.244825 ],
  22.         [-1.2101456]], dtype=float32)>,
  23.  <tf.Variable 'dense_2/bias:0' shape=(1,) dtype=float32, numpy=array([0.], dtype=float32)>]
  1. model.layers[0].trainable = False #冻结第0层的变量,使其不可训练
  2. model.trainable_variables
  1. [<tf.Variable 'dense_1/kernel:0' shape=(4, 2) dtype=float32, numpy=
  2.  array([[ 0.5022683 , -0.0507431 ],
  3.         [-0.61540484,  0.9369011 ],
  4.         [-0.14412141, -0.54607415],
  5.         [ 0.2027781 , -0.4651153 ]], dtype=float32)>,
  6.  <tf.Variable 'dense_1/bias:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>,
  7.  <tf.Variable 'dense_2/kernel:0' shape=(2, 1) dtype=float32, numpy=
  8.  array([[-0.244825 ],
  9.         [-1.2101456]], dtype=float32)>,
  10.  <tf.Variable 'dense_2/bias:0' shape=(1,) dtype=float32, numpy=array([0.], dtype=float32)>]
  1. model.submodules
  1. (<tensorflow.python.keras.engine.input_layer.InputLayer at 0x144d8c080>,
  2.  <tensorflow.python.keras.layers.core.Dense at 0x144daada0>,
  3.  <tensorflow.python.keras.layers.core.Dense at 0x144d8c5c0>,
  4.  <tensorflow.python.keras.layers.core.Dense at 0x144d7aa20>)
  1. model.layers
  1. [<tensorflow.python.keras.layers.core.Dense at 0x144daada0>,
  2.  <tensorflow.python.keras.layers.core.Dense at 0x144d8c5c0>,
  3.  <tensorflow.python.keras.layers.core.Dense at 0x144d7aa20>]
  1. print(model.name)
  2. print(model.name_scope())
  1. sequential
  2. sequential

如果对本书内容理解上有需要进一步和作者交流的地方，欢迎在公众号”Python与算法之美”下留言。作者时间和精力有限，会酌情予以回复。

也可以在公众号后台回复关键字：加群，加入读者交流群和大家讨论。

image.png