AutoGraph和tf.Module
1.Autograph和tf.Module概述
构建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开发任意机器学习模型(而非仅仅是神经网络模型),并实现跨平台部署使用。
2.应用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)) #4
利用tf.Module的子