本文是对官方文档 的学习笔记。
抽象的说, TF2 模型是：
- 在Tensor 上运行的计算
- 可以根据训练更新的 Variable
什么是模型（models）和层（layers）
绝大多数模型都由层组成。层是有某种数学结构的函数，该函数可以被重用并且拥有可以被训练的variable。 再 TF2 中， high-level对layers 和 model 的实现，例如 Keras 和 Sonnet 都基于基础类 tf.Module.
一个Model 例子
class SimpleModule(tf.Module):
def __init__(self, name=None):
super().__init__(name=name)
self.a_variable = tf.Variable(5.0, name="train_me")
self.non_trainable_variable = tf.Variable(5.0, trainable=False, name="do_not_train_me")
def __call__(self, x):
return self.a_variable * x + self.non_trainable_variable
simple_module = SimpleModule(name="simple")
simple_module(tf.constant(5.0))
-----------------------------------------------------------------------
<tf.Tensor: shape=(), dtype=float32, numpy=30.0>
Module 和 Layer 都是深度学习里的术语， 它们有内部状态， 有利用那些内部状态的方法。
注意： tf.Module 是 tf.keras.layers.Layer 和 tf.keras.Model 的基类。 处于兼容的考虑， Keras Layer 不会用 Module 来管理 Variable，所以 Model 要么使用 Module 要么使用 Layers。
通过继承 tf.Module， 任何赋值给该对象的 tf.Variable 或者 tf.Model 实例都会被收集。 从而，可以保持，加载 variable, 并且创建 tf.Module 的集合。
# All trainable variables
print("trainable variables:", simple_module.trainable_variables)
# Every variable
print("all variables:", simple_module.variables)
----------------------------------------------------------
trainable variables: (<tf.Variable 'train_me:0' shape=() dtype=float32, numpy=5.0>,)
all variables: (<tf.Variable 'train_me:0' shape=() dtype=float32, numpy=5.0>, <tf.Variable 'do_not_train_me:0' shape=() dtype=float32, numpy=5.0>)
下面是一个用module 做成的 双层线性 layer 模型。
首先是一个 Dense Layer
class Dense(tf.Module):
def __init__(self, in_features, out_features, name=None):
super().__init__(name=name)
self.w = tf.Variable(
tf.random.normal([in_features, out_features]), name='w')
self.b = tf.Variable(tf.zeros([out_features]), name='b')
def __call__(self, x):
y = tf.matmul(x, self.w) + self.b
return tf.nn.relu(y)
双层模型
class SequentialModule(tf.Module):
def __init__(self, name=None):
super().__init__(name=name)
self.dense_1 = Dense(in_features=3, out_features=3)
self.dense_2 = Dense(in_features=3, out_features=2)
def __call__(self, x):
x = self.dense_1(x)
return self.dense_2(x)
# You have made a model!
my_model = SequentialModule(name="the_model")
# Call it, with random results
print("Model results:", my_model(tf.constant([[2.0, 2.0, 2.0]])))
对上面这个简单模型的调用
class SequentialModule(tf.Module):
def __init__(self, name=None):
super().__init__(name=name)
self.dense_1 = Dense(in_features=3, out_features=3)
self.dense_2 = Dense(in_features=3, out_features=2)
def __call__(self, x):
x = self.dense_1(x)
return self.dense_2(x)
# You have made a model!
my_model = SequentialModule(name="the_model")
# Call it, with random results
print("Model results:", my_model(tf.constant([[2.0, 2.0, 2.0]])))
------------------------------------------------------
Model results: tf.Tensor([[0. 1.0405664]], shape=(1, 2), dtype=float32)
tf.Module 实例会自动管理赋值给它 tf.Variable / tf.Module 实例。 这样可以大大简化相关操作。
print("Submodules:", my_model.submodules)
--------------------------------------------------------------
Submodules: (<__main__.Dense object at 0x7f8c5d4816a0>, <__main__.Dense object at 0x7f8c5d481710>)
for var in my_model.variables:
print(var, "\n")
--------------------------------------------------------------
<tf.Variable 'b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>
<tf.Variable 'w:0' shape=(3, 3) dtype=float32, numpy=
array([[ 1.3512987 , -1.0520874 , 0.01668975],
[-0.23172653, -0.29609755, 0.23737088],
[ 1.4239831 , 1.3281558 , -0.5324163 ]], dtype=float32)>
<tf.Variable 'b:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>
<tf.Variable 'w:0' shape=(3, 2) dtype=float32, numpy=
array([[-0.91189057, 0.2045496 ],
[-0.01225489, -0.02003982],
[-0.32654697, -0.501168 ]], dtype=float32)>
上面例子比较神奇的部分在于： self.w, self.b 都是普通的类成员， 但是一旦 tf.Variable 实例赋值给他， 他就能发现这些 variable， 并用 self.variables 变量引用到这些实例。
延迟创建变量
上例中， 一开始给出了Input ， outpus Tensor 的shape， 因此w变量具有已知的形状并且可以分配。这不是必须的， 也可以等到第一次实例化时在确定shape，这样就无需预先指定输入大小。
class FlexibleDenseModule(tf.Module):
# Note: No need for `in+features`
def __init__(self, out_features, name=None):
super().__init__(name=name)
self.is_built = False
self.out_features = out_features
def __call__(self, x):
# Create variables on first call.
if not self.is_built:
self.w = tf.Variable(
tf.random.normal([x.shape[-1], self.out_features]), name='w')
self.b = tf.Variable(tf.zeros([self.out_features]), name='b')
self.is_built = True
y = tf.matmul(x, self.w) + self.b
return tf.nn.relu(y)
# Used in a module
class MySequentialModule(tf.Module):
def __init__(self, name=None):
super().__init__(name=name)
self.dense_1 = FlexibleDenseModule(out_features=3)
self.dense_2 = FlexibleDenseModule(out_features=2)
def __call__(self, x):
x = self.dense_1(x)
return self.dense_2(x)
my_model = MySequentialModule(name="the_model")
print("Model results:", my_model(tf.constant([[2.0, 2.0, 2.0]])))
----------------------------------------------------------------------
Model results: tf.Tensor([[0.9369497 0. ]], shape=(1, 2), dtype=float32)
这使得 TensorFlow layer 通常只需要指定其输出形状的原因，例如在tf.keras.layers.Dense中，而不是输入和输出大小。
保存权重
可以用 checkpoint 或者 SavedModel 来保存 Module。
检查点只是权重（即模块及其子模块内部的变量集的值）：
chkp_path = "my_checkpoint"
checkpoint = tf.train.Checkpoint(model=my_model)
checkpoint.write(chkp_path)
---------------------------------------------
'my_checkpoint'
检查点由两种文件组成：数据和元数据的索引文件。索引文件跟踪实际保存的内容和检查点的编号，而 checkpoint 数据包含变量值及其属性查找路径
ls my_checkpoint*
my_checkpoint.data-00000-of-00001 my_checkpoint.index
查看一个检查点，以确保保存了所有变量集，并按包含变量的Python对象对其进行了排序。
tf.train.list_variables(chkp_path)
---------------------------------------------
[('_CHECKPOINTABLE_OBJECT_GRAPH', []),
('model/dense_1/b/.ATTRIBUTES/VARIABLE_VALUE', [3]),
('model/dense_1/w/.ATTRIBUTES/VARIABLE_VALUE', [3, 3]),
('model/dense_2/b/.ATTRIBUTES/VARIABLE_VALUE', [2]),
('model/dense_2/w/.ATTRIBUTES/VARIABLE_VALUE', [3, 2])]
在分布式（多机）训练期间，可以将它们分片，这就是为什么要对它们进行编号的原因（例如'00000-of-00001'）。但是，在这种情况下，只有一个分片。
重新加载模型时，将覆盖Python对象中的值。
new_model = MySequentialModule()
new_checkpoint = tf.train.Checkpoint(model=new_model)
new_checkpoint.restore("my_checkpoint")
# Should be the same result as above
new_model(tf.constant([[2.0, 2.0, 2.0]]))
--------------------------------------------------------------
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[0.9369497, 0. ]], dtype=float32)>
注意：由于checkpoint 是training 流程的核心，因此tf.checkpoint.CheckpointManager是一个帮助程序类，它使检查点的管理更加容易。有关更多详细信息，请参 Training checkpoints guide
存储函数
TensorFlow可以在没有原始Python对象的情况下运行模型，正如TensorFlow Serving和TensorFlow Lite所展示的，即使从TensorFlow Hub下载经过训练的模型也是如此。 TensorFlow需要知道如何执行Python中描述的计算，即使没有原代码。为此，您可以制作一个 graph，在 Introduction to graphs and functions guide。
可以通过添加@ tf.function装饰器在上面的模型中定义图，以指示该代码应作为图运行。
class MySequentialModule(tf.Module):
def __init__(self, name=None):
super().__init__(name=name)
self.dense_1 = Dense(in_features=3, out_features=3)
self.dense_2 = Dense(in_features=3, out_features=2)
@tf.function
def __call__(self, x):
x = self.dense_1(x)
return self.dense_2(x)
# You have made a model with a graph!
my_model = MySequentialModule(name="the_model")
print(my_model([[2.0, 2.0, 2.0]]))
print(my_model([[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]]))
-----------------------------------------------------------------
tf.Tensor([[8.405768 0. ]], shape=(1, 2), dtype=float32)
tf.Tensor(
[[[8.405768 0. ]
[8.405768 0. ]]], shape=(1, 2, 2), dtype=float32)
用 TensorBoard 来看构建的 graph
# Set up logging.
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = "logs/func/%s" % stamp
writer = tf.summary.create_file_writer(logdir)
# Create a new model to get a fresh trace
# Otherwise the summary will not see the graph.
new_model = MySequentialModule()
# Bracket the function call with
# tf.summary.trace_on() and tf.summary.trace_export().
tf.summary.trace_on(graph=True)
tf.profiler.experimental.start(logdir)
# Call only one tf.function when tracing.
z = print(new_model(tf.constant([[2.0, 2.0, 2.0]])))
with writer.as_default():
tf.summary.trace_export(
name="my_func_trace",
step=0,
profiler_outdir=logdir)
---------------------------------------------------------------------------
tf.Tensor([[0. 0.2614944]], shape=(1, 2), dtype=float32)
===========================================
%tensorboard --logdir logs/func
image.png
创建 SavedModel
推荐使用SavedModel。 SavedModel包含函数集合和权重集合。
tf.saved_model.save(my_model, "the_saved_model")
---------------------------------------------------------------------------
INFO:tensorflow:Assets written to: the_saved_model/assets
# Inspect the SavedModel in the directory
ls -l the_saved_model
---------------------------------------------------------------------------
total 24
drwxr-sr-x 2 kbuilder kokoro 4096 Jan 13 02:26 assets
-rw-rw-r-- 1 kbuilder kokoro 14144 Jan 13 02:26 saved_model.pb
drwxr-sr-x 2 kbuilder kokoro 4096 Jan 13 02:26 variables
# The variables/ directory contains a checkpoint of the variables
ls -l the_saved_model/variables
---------------------------------------------------------------------------
total 8
-rw-rw-r-- 1 kbuilder kokoro 408 Jan 13 02:26 variables.data-00000-of-00001
-rw-rw-r-- 1 kbuilder kokoro 356 Jan 13 02:26 variables.index
saved_model.pb 文件是描述 functional tf.Graph 的 protocol buffer 文件。
可以从此表示形式加载模型和图层，而无需实际创建创建它的类的实例。在没有（或不需要）Python解释器的情况下（例如大规模或在边缘设备上使用），或者在原始Python代码不可用或不实际使用的情况下，这是理想的。
new_model = tf.saved_model.load("the_saved_model")
通过加载保存的模型创建的new_model是内部TensorFlow用户对象，无需任何类知识。它不是SequentialModule类型的。
isinstance(new_model, SequentialModule)
---------------------------------------------------------------
False
此新模型适用于已定义的输入签名。不能向这样还原的模型添加更多签名。
print(my_model([[2.0, 2.0, 2.0]]))
print(my_model([[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]]))
---------------------------------------------------------------
tf.Tensor([[8.405768 0. ]], shape=(1, 2), dtype=float32)
tf.Tensor(
[[[8.405768 0. ]
[8.405768 0. ]]], shape=(1, 2, 2), dtype=float32)
因此，使用SavedModel，您可以使用tf.Module保存TensorFlow权重和图形，然后再次加载它们。
Keras models and layers
上面介绍的是 Moduel ， 接下来介绍 Keras 的Model 类和 Layer 如何利用 module 类。
Keras layer
tf.keras.layers.Layer 是 Keras 所有layer 的基类， 而它本身是 Module 的子类。
只需要更换父类，然后将call更改为call即可将 Module 转换为Keras layer：
class MyDense(tf.keras.layers.Layer):
# Adding **kwargs to support base Keras layer arguments
def __init__(self, in_features, out_features, **kwargs):
super().__init__(**kwargs)
# This will soon move to the build step; see below
self.w = tf.Variable(
tf.random.normal([in_features, out_features]), name='w')
self.b = tf.Variable(tf.zeros([out_features]), name='b')
def call(self, x):
y = tf.matmul(x, self.w) + self.b
return tf.nn.relu(y)
simple_layer = MyDense(name="simple", in_features=3, out_features=3)
Keras layer 具有自己的call，它会进行 bookkeeping ，然后调用call（）。在功能上并没有差别。
simple_layer([[2.0, 2.0, 2.0]])
--------------------------------------------------------------
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=array([[6.8031726, 0. , 0. ]], dtype=float32)>
Build
如上所述，确定 input shape 之前不变量在许多情况下都很方便。
Keras图层附带了一个额外的生命周期步骤（build），可以更加灵活地定义 layer。这在构建函数中定义。 build仅被调用一次，并以输入的 shape 进行调用。通常用于创建变量（weights）。
可以重写上面的MyDense层以灵活地调整其输入的大小：
class FlexibleDense(tf.keras.layers.Layer):
# Note the added `**kwargs`, as Keras supports many arguments
def __init__(self, out_features, **kwargs):
super().__init__(**kwargs)
self.out_features = out_features
def build(self, input_shape): # Create the state of the layer (weights)
self.w = tf.Variable(
tf.random.normal([input_shape[-1], self.out_features]), name='w')
self.b = tf.Variable(tf.zeros([self.out_features]), name='b')
def call(self, inputs): # Defines the computation from inputs to outputs
return tf.matmul(inputs, self.w) + self.b
# Create the instance of the layer
flexible_dense = FlexibleDense(out_features=3)
这时，model build 还没有被调用， 所以没有 variables
flexible_dense.variables
-----------------------------------------------
[]
调用该函数会分配适当大小的变量：
# Call it, with predictably random results
print("Model results:", flexible_dense(tf.constant([[2.0, 2.0, 2.0], [3.0, 3.0, 3.0]])))
-----------------------------------------------
Model results: tf.Tensor(
[[-0.19694364 -0.16952538 -0.8057083 ]
[-0.2954154 -0.2542882 -1.2085624 ]], shape=(2, 3), dtype=float32)
flexible_dense.variables
-----------------------------------------------
[<tf.Variable 'flexible_dense/w:0' shape=(3, 3) dtype=float32, numpy=
array([[ 0.17480798, -0.77812535, -0.86842895],
[-1.0463405 , 0.3601503 , -0.4638048 ],
[ 0.7730607 , 0.33321235, 0.92937964]], dtype=float32)>,
<tf.Variable 'flexible_dense/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>]
由于仅调用一次build，因此如果 input shape 与 layer variable 不兼容，则 input 将被拒绝：
try:
print("Model results:", flexible_dense(tf.constant([[2.0, 2.0, 2.0, 2.0]])))
except tf.errors.InvalidArgumentError as e:
print("Failed:", e)
-----------------------------------------------
Failed: Matrix size-incompatible: In[0]: [1,4], In[1]: [3,3] [Op:MatMul]
Keras layer 还有一些其他的特点：
- 可选的 Optional losses
- 支持 metrics
- 内置对可选训练参数的支持，以区分训练和预测使用
- get_config 和 from_config 方法，以支持准确的保存，加载和 clone。
Keras model
可以将模型定义为嵌套的Keras图层。
但是，Keras还提供了名为tf.keras.Model的全功能模型类。它继承自tf.keras.layers.Layer，因此可以使用与Keras图层相同的方式使用，嵌套和保存Keras模型。 Keras模型具有额外的功能，使它们易于在多台机器上进行训练，评估，加载，保存甚至训练。
class MySequentialModel(tf.keras.Model):
def __init__(self, name=None, **kwargs):
super().__init__(**kwargs)
self.dense_1 = FlexibleDense(out_features=3)
self.dense_2 = FlexibleDense(out_features=2)
def call(self, x):
x = self.dense_1(x)
return self.dense_2(x)
# You have made a Keras model!
my_sequential_model = MySequentialModel(name="the_model")
# Call it on a tensor, with random results
print("Model results:", my_sequential_model(tf.constant([[2.0, 2.0, 2.0]])))
-----------------------------------------------
Model results: tf.Tensor([[0.35962012 3.3997526 ]], shape=(1, 2), dtype=float32)
所有相同的功能都可用，包括tracking variables和submodules。
my_sequential_model.variables
-----------------------------------------------
[<tf.Variable 'my_sequential_model/flexible_dense_1/w:0' shape=(3, 3) dtype=float32, numpy=
array([[ 0.28385058, 0.7153611 , -0.7905529 ],
[-0.23553997, -1.0608387 , -1.9260269 ],
[-0.5398518 , 0.6666878 , -0.38458622]], dtype=float32)>,
<tf.Variable 'my_sequential_model/flexible_dense_1/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>,
<tf.Variable 'my_sequential_model/flexible_dense_2/w:0' shape=(3, 2) dtype=float32, numpy=
array([[ 2.1480818 , -0.7663207 ],
[ 0.36995146, 0.15208478],
[-0.36013824, -0.41092506]], dtype=float32)>,
<tf.Variable 'my_sequential_model/flexible_dense_2/b:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>]
my_sequential_model.submodules
-----------------------------------------------
(<__main__.FlexibleDense at 0x7f8cc7cfbf60>,
<__main__.FlexibleDense at 0x7f8cc7cb70f0>)
覆盖tf.keras.Model是构建TensorFlow模型的非常Python化的方法。如果要从其他框架迁移模型，这可能非常简单。
如果要构建的模型是现有图层和输入的简单组合，则可以使用功能性API节省时间和空间，该功能性API附带了有关模型重构和体系结构的其他功能。
用 API 方法构建起来的模型
inputs = tf.keras.Input(shape=[3,])
x = FlexibleDense(3)(inputs)
x = FlexibleDense(2)(x)
my_functional_model = tf.keras.Model(inputs=inputs, outputs=x)
my_functional_model.summary()
-----------------------------------------------
Model: "model"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_1 (InputLayer) [(None, 3)] 0
_________________________________________________________________
flexible_dense_3 (FlexibleDe (None, 3) 12
_________________________________________________________________
flexible_dense_4 (FlexibleDe (None, 2) 8
=================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
my_functional_model(tf.constant([[2.0, 2.0, 2.0]]))
--------------------------------------------------------------
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[-4.795524 , -0.43587577]], dtype=float32)>
这里的主要区别在于，输入形状是功能构建过程的一部分，是预先指定的。在这种情况下，不必完全指定input_shape参数。您可以将某些尺寸保留为“无”。
保存 Keras 模型
可以对Keras模型进行检查，其外观将与tf.Module相同。
Keras模型也可以使用tf.saved_models.save（）保存，因为它们是模块。但是，Keras模型具有便捷方法和其他功能：
my_sequential_model.save("exname_of_file")
--------------------------------------------------------------
INFO:tensorflow:Assets written to: exname_of_file/assets
同样容易地，它们可以重新载入：
reconstructed_model = tf.keras.models.load_model("exname_of_file")
--------------------------------------------------------------
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
Keras SavedModels还保存merics ，loss 和 optimizer 状态。
可以使用此重构的模型，并在对相同数据进行调用时将产生相同的结果：
reconstructed_model(tf.constant([[2.0, 2.0, 2.0]]))
--------------------------------------------------------------
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[0.35962012, 3.3997526 ]], dtype=float32)>
有关Keras模型的保存和序列化的更多信息，包括为自定义图层提供配置方法以支持功能。查看guide to saving and serialization。
网友评论