 在 TensorFlow.org 上查看
|
 在 Google Colab 中运行
|
 在 GitHub 上查看源代码
|
 下载笔记本
|
要在 TensorFlow 中进行机器学习，您可能需要定义、保存和恢复模型。
抽象地说，模型是
- 一个在张量上计算某些内容的函数（**前向传递**）
- 一些可以根据训练进行更新的变量
在本指南中，您将深入了解 Keras 的表面，了解 TensorFlow 模型的定义方式。这将介绍 TensorFlow 如何收集变量和模型，以及如何保存和恢复它们。
设置
import tensorflow as tf
import keras
from datetime import datetime
%load_ext tensorboard
2023-10-18 01:21:05.536666: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2023-10-18 01:21:05.536712: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2023-10-18 01:21:05.536766: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
TensorFlow 模块
大多数模型都是由层组成的。层是具有已知数学结构的函数，可以重复使用并具有可训练的变量。在 TensorFlow 中，大多数层和模型的高级实现（如 Keras 或 Sonnet）都是基于同一个基础类构建的：tf.Module。
构建模块
以下是一个非常简单的 tf.Module 示例，它对标量张量进行操作
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))
2023-10-18 01:21:08.181350: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2211] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://tensorflowcn.cn/install/gpu for how to download and setup the required libraries for your platform. Skipping registering GPU devices... <tf.Tensor: shape=(), dtype=float32, numpy=30.0>
模块以及扩展的层是深度学习术语中的“对象”：它们具有内部状态，以及使用该状态的方法。
关于 __call__ 没有什么特别之处，除了像 Python 可调用对象 一样；您可以使用您想要的任何函数来调用您的模型。
您可以出于任何原因打开和关闭变量的可训练性，包括在微调期间冻结层和变量。
通过子类化 tf.Module，分配给此对象属性的任何 tf.Variable 或 tf.Module 实例都会被自动收集。这允许您保存和加载变量，以及创建 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>)
这是一个由模块组成的两层线性层模型的示例。
首先是一个密集（线性）层
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]])))
Model results: tf.Tensor([[0. 3.415034]], shape=(1, 2), dtype=float32)
tf.Module 实例会自动递归地收集分配给它的任何 tf.Variable 或 tf.Module 实例。这允许您使用单个模型实例管理 tf.Module 的集合，并保存和加载整个模型。
print("Submodules:", my_model.submodules)
Submodules: (<__main__.Dense object at 0x7f7931aea250>, <__main__.Dense object at 0x7f77ed5b8a00>)
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([[-2.8161757, -2.6065955, 1.9061812],
[-0.9430401, -0.4624743, -0.4531979],
[-1.3428234, 0.7062293, 0.7874674]], 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([[ 1.0474309 , -0.6920227 ],
[ 1.2405277 , 0.36411622],
[-1.6990206 , 0.762131 ]], dtype=float32)>
等待创建变量
您可能已经注意到，您必须为层定义输入和输出大小。这是为了使 w 变量具有已知的形状并可以分配。
通过将变量创建推迟到模块第一次使用特定输入形状调用时，您无需预先指定输入大小。
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. 0.]], shape=(1, 2), dtype=float32)
这种灵活性是 TensorFlow 层通常只需要指定其输出的形状的原因，例如在 tf.keras.layers.Dense 中，而不是同时指定输入和输出大小。
保存权重
您可以将 tf.Module 同时保存为 检查点 和 SavedModel。
检查点只是权重（即模块及其子模块内部变量集的值）
chkp_path = "my_checkpoint"
checkpoint = tf.train.Checkpoint(model=my_model)
checkpoint.write(chkp_path)
'my_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., 0.]], dtype=float32)>
保存函数
TensorFlow 可以运行没有原始 Python 对象的模型，如 TensorFlow Serving 和 TensorFlow Lite 所示，即使您从 TensorFlow Hub 下载了训练过的模型。
TensorFlow 需要知道如何执行 Python 中描述的计算，但不需要原始代码。为此，您可以创建一个图，这在 图和函数介绍指南 中有描述。
此图包含实现该函数的操作或运算。
您可以在上面的模型中通过添加 @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([[0.31593648 0. ]], shape=(1, 2), dtype=float32) tf.Tensor( [[[0.31593648 0. ] [0.31593648 0. ]]], shape=(1, 2, 2), dtype=float32)
您可以通过在 TensorBoard 摘要中跟踪图来可视化它。
# 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.]], shape=(1, 2), dtype=float32)
启动 TensorBoard 以查看生成的跟踪
#docs_infra: no_execute
%tensorboard --logdir logs/func
创建 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 directoryls -l the_saved_model
total 32 drwxr-sr-x 2 kbuilder kokoro 4096 Oct 18 01:21 assets -rw-rw-r-- 1 kbuilder kokoro 58 Oct 18 01:21 fingerprint.pb -rw-rw-r-- 1 kbuilder kokoro 17704 Oct 18 01:21 saved_model.pb drwxr-sr-x 2 kbuilder kokoro 4096 Oct 18 01:21 variables
# The variables/ directory contains a checkpoint of the variablesls -l the_saved_model/variables
total 8 -rw-rw-r-- 1 kbuilder kokoro 490 Oct 18 01:21 variables.data-00000-of-00001 -rw-rw-r-- 1 kbuilder kokoro 356 Oct 18 01:21 variables.index
saved_model.pb 文件是一个 协议缓冲区，描述了功能性的 tf.Graph。
模型和层可以从这种表示中加载，而无需实际创建创建它的类的实例。这在您没有（或不需要）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([[0.31593648 0. ]], shape=(1, 2), dtype=float32) tf.Tensor( [[[0.31593648 0. ] [0.31593648 0. ]]], shape=(1, 2, 2), dtype=float32)
因此，使用 SavedModel，您可以使用 tf.Module 保存 TensorFlow 权重和图，然后再次加载它们。
Keras 模型和层
请注意，到目前为止，还没有提到 Keras。您可以在 tf.Module 之上构建自己的高级 API，而且人们已经这样做了。
在本节中，您将研究 Keras 如何使用 tf.Module。Keras 模型的完整用户指南可以在 Keras 指南 中找到。
Keras 层和模型具有更多额外功能，包括
- 可选损失
- 支持 指标
- 内置支持可选的
training参数，以区分训练和推理使用 - 保存和恢复 Python 对象，而不仅仅是黑盒函数
get_config和from_config方法，允许您准确地存储配置以允许在 Python 中克隆模型
这些功能允许通过子类化创建更复杂的模型，例如自定义 GAN 或变分自动编码器 (VAE) 模型。在 自定义层和模型完整指南 中阅读有关它们的更多信息。
Keras 模型还附带了额外的功能，使它们易于训练、评估、加载、保存，甚至在多台机器上训练。
Keras 层
tf.keras.layers.Layer 是所有 Keras 层的基类，它继承自 tf.Module。
您可以通过交换父类，然后将 __call__ 更改为 call 来将模块转换为 Keras 层
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 层有自己的 __call__，它会执行下一节中描述的一些簿记工作，然后调用 call()。您应该注意到功能没有变化。
simple_layer([[2.0, 2.0, 2.0]])
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=array([[1.1688161, 0. , 0. ]], dtype=float32)>
build 步骤
如前所述，在许多情况下，将变量创建推迟到确定输入形状后是方便的。
Keras 层附带了一个额外的生命周期步骤，允许您在定义层的方式上更加灵活。这在 build 函数中定义。
build 恰好被调用一次，并且它使用输入的形状被调用。它通常用于创建变量（权重）。
您可以将上面的 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)
此时，模型尚未构建，因此没有变量
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( [[-2.531786 -5.5550847 -0.4248762] [-3.7976792 -8.332626 -0.6373143]], shape=(2, 3), dtype=float32)
flexible_dense.variables
[<tf.Variable 'flexible_dense/w:0' shape=(3, 3) dtype=float32, numpy=
array([[-0.77719826, -1.9281565 , 0.82326293],
[ 0.85628736, -0.31845194, 0.10916236],
[-1.3449821 , -0.5309338 , -1.1448634 ]], dtype=float32)>,
<tf.Variable 'flexible_dense/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>]
由于 build 仅被调用一次，因此如果输入形状与层的变量不兼容，则会拒绝输入
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: Exception encountered when calling layer 'flexible_dense' (type FlexibleDense).
{ {function_node __wrapped__MatMul_device_/job:localhost/replica:0/task:0/device:CPU:0} } Matrix size-incompatible: In[0]: [1,4], In[1]: [3,3] [Op:MatMul] name:
Call arguments received by layer 'flexible_dense' (type FlexibleDense):
• inputs=tf.Tensor(shape=(1, 4), dtype=float32)
Keras 模型
您可以将模型定义为嵌套的 Keras 层。
但是，Keras 还提供了一个功能齐全的模型类，称为 tf.keras.Model。它继承自 tf.keras.layers.Layer，因此 Keras 模型可以使用和嵌套的方式与 Keras 层相同。Keras 模型附带了额外的功能，使它们易于训练、评估、加载、保存，甚至在多台机器上训练。
您可以使用几乎相同的代码定义上面的 SequentialModule，再次将 __call__ 转换为 call() 并更改父类
@keras.saving.register_keras_serializable()
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.26034355 16.431221 ]], shape=(1, 2), dtype=float32)
所有相同的功能都可用，包括跟踪变量和子模块。
my_sequential_model.variables
[<tf.Variable 'my_sequential_model/flexible_dense_1/w:0' shape=(3, 3) dtype=float32, numpy=
array([[ 1.4749854 , 0.16090827, 2.2669017 ],
[ 1.6850946 , 1.1545411 , 0.1707306 ],
[ 0.8753734 , -0.13549292, 0.08751986]], 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([[-0.8022977 , 1.9773549 ],
[-0.76657015, -0.8485579 ],
[ 1.6919082 , 0.49000967]], 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 0x7f790c7e0e80>, <__main__.FlexibleDense at 0x7f790c7e6940>)
覆盖 tf.keras.Model 是构建 TensorFlow 模型的一种非常 Pythonic 的方法。如果您正在从其他框架迁移模型，这将非常简单。
如果您正在构建简单的现有层和输入组合的模型，您可以通过使用 函数式 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 (Flexible (None, 3) 12
Dense)
flexible_dense_4 (Flexible (None, 2) 8
Dense)
=================================================================
Total params: 20 (80.00 Byte)
Trainable params: 20 (80.00 Byte)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________
my_functional_model(tf.constant([[2.0, 2.0, 2.0]]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[3.4276495, 2.937252 ]], dtype=float32)>
这里的主要区别在于输入形状是在函数式构建过程的开始就指定的。在这种情况下，input_shape 参数不必完全指定；您可以将某些维度保留为 None。
保存 Keras 模型
Keras 模型有自己的专用 zip 存档保存格式，以 .keras 扩展名标记。当调用 tf.keras.Model.save 时，请在文件名中添加 .keras 扩展名。例如
my_sequential_model.save("exname_of_file.keras")
同样容易，它们可以被加载回来
reconstructed_model = tf.keras.models.load_model("exname_of_file.keras")
Keras zip 存档（.keras 文件）还保存指标、损失和优化器状态。
此重建的模型可以使用，并且在对相同数据调用时会产生相同的结果
reconstructed_model(tf.constant([[2.0, 2.0, 2.0]]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[ 0.26034355, 16.431221 ]], dtype=float32)>
检查 Keras 模型
Keras 模型也可以进行检查点，这与 tf.Module 相似。
关于 Keras 模型的保存和序列化还有更多内容，包括为自定义层提供配置方法以支持功能。请查看 保存和序列化的指南。
下一步
如果您想了解更多关于 Keras 的详细信息，您可以查看现有的 Keras 指南 这里。
另一个基于 tf.module 构建的高级 API 示例是来自 DeepMind 的 Sonnet，它在 他们的网站 上有介绍。