4-5,AutoGraph和tf.Module

    TensorFlow 2.0主要使用的是动态计算图和Autograph。

    动态计算图易于调试,编码效率较高,但执行效率偏低。

    静态计算图执行效率很高,但较难调试。

    而Autograph机制可以将动态图转换成静态计算图,兼收执行效率和编码效率之利。

    当然Autograph机制能够转换的代码并不是没有任何约束的,有一些编码规范需要遵循,否则可能会转换失败或者不符合预期。

    前面我们介绍了Autograph的编码规范和Autograph转换成静态图的原理。

    本篇我们介绍使用tf.Module来更好地构建Autograph。

    但是如果在函数外部定义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开发任意机器学习模型(而非仅仅是神经网络模型),并实现跨平台部署使用。

    定义一个简单的function。

    1. #add_print(tf.constant(3)) #输入不符合张量签名的参数将报错
    1. 4
    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)
    6. @tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.float32)])
    7. def addprint(self,a):
    8. with self.name_scope:
    9. self.x.assign_add(a)
    10. tf.print(self.x)
    11. 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

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

    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.x = tf.Variable(0.0)
    2. def addprint(a):
    3. mymodule.x.assign_add(a)
    4. tf.print(mymodule.x)
    5. return (mymodule.x)
    6. 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})
    4. #加载模型
    5. mymodule2 = tf.saved_model.load("./data/mymodule")
    6. mymodule2.addprint(tf.constant(5.0))
    1. INFO:tensorflow:Assets written to: ./data/mymodule/assets
    2. 5

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

    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()
    2. model = models.Sequential()
    3. model.add(layers.Dense(4,input_shape = (10,)))
    4. model.add(layers.Dense(2))
    5. model.add(layers.Dense(1))
    6. 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. Non-trainable params: 0
    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())

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

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