<jupyter_start><jupyter_text>Copyright 2019 The TensorFlow Authors.<jupyter_code>#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.<jupyter_output><empty_output><jupyter_text>使用分布策略保存和加载模型   在 TensorFlow.org 上查看  在 Google Colab 中运行  在 Github 上查看源代码  下载笔记本  概述在训练期间一般需要保存和加载模型。有两组用于保存和加载 Keras 模型的 API：高级 API 和低级 API。本教程演示了在使用 `tf.distribute.Strategy` 时如何使用 SavedModel API。要了解 SavedModel 和序列化的相关概况，请参阅[保存的模型指南](../../guide/saved_model.ipynb)和 [Keras 模型序列化指南](../../guide/keras/save_and_serialize.ipynb)。让我们从一个简单的示例开始：  导入依赖项：<jupyter_code>import tensorflow_datasets as tfds

import tensorflow as tf
tfds.disable_progress_bar()<jupyter_output><empty_output><jupyter_text>使用 `tf.distribute.Strategy` 准备数据和模型：<jupyter_code>mirrored_strategy = tf.distribute.MirroredStrategy()

def get_data():
  datasets, ds_info = tfds.load(name='mnist', with_info=True, as_supervised=True)
  mnist_train, mnist_test = datasets['train'], datasets['test']

  BUFFER_SIZE = 10000

  BATCH_SIZE_PER_REPLICA = 64
  BATCH_SIZE = BATCH_SIZE_PER_REPLICA * mirrored_strategy.num_replicas_in_sync

  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255

    return image, label

  train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
  eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)

  return train_dataset, eval_dataset

def get_model():
  with mirrored_strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
        tf.keras.layers.MaxPooling2D(),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10)
    ])

    model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  optimizer=tf.keras.optimizers.Adam(),
                  metrics=['accuracy'])
    return model<jupyter_output><empty_output><jupyter_text>训练模型：<jupyter_code>model = get_model()
train_dataset, eval_dataset = get_data()
model.fit(train_dataset, epochs=2)<jupyter_output><empty_output><jupyter_text>保存和加载模型现在，您已经有一个简单的模型可供使用，让我们了解一下如何保存/加载 API。有两组可用的 API：- 高级 Keras `model.save` 和 `tf.keras.models.load_model`- 低级 `tf.saved_model.save` 和 `tf.saved_model.load`  Keras API 以下为使用 Keras API 保存和加载模型的示例：<jupyter_code>keras_model_path = "/tmp/keras_save"
model.save(keras_model_path)  # save() should be called out of strategy scope<jupyter_output><empty_output><jupyter_text>恢复无 `tf.distribute.Strategy` 的模型：<jupyter_code>restored_keras_model = tf.keras.models.load_model(keras_model_path)
restored_keras_model.fit(train_dataset, epochs=2)<jupyter_output><empty_output><jupyter_text>恢复模型后，您可以继续在它上面进行训练，甚至无需再次调用 `compile()`，因为在保存之前已经对其进行了编译。模型以 TensorFlow 的标准 `SavedModel` proto 格式保存。有关更多信息，请参阅 [`saved_model` 格式指南](../../guide/saved_model.ipynb)。现在，加载模型并使用 `tf.distribute.Strategy` 进行训练：<jupyter_code>another_strategy = tf.distribute.OneDeviceStrategy("/cpu:0")
with another_strategy.scope():
  restored_keras_model_ds = tf.keras.models.load_model(keras_model_path)
  restored_keras_model_ds.fit(train_dataset, epochs=2)<jupyter_output><empty_output><jupyter_text>如您所见， `tf.distribute.Strategy` 可以按预期进行加载。此处使用的策略不必与保存前所用策略相同。   `tf.saved_model` API 现在，让我们看一下较低级别的 API。保存模型与 Keras API 类似：<jupyter_code>model = get_model()  # get a fresh model
saved_model_path = "/tmp/tf_save"
tf.saved_model.save(model, saved_model_path)<jupyter_output><empty_output><jupyter_text>可以使用 `tf.saved_model.load()` 进行加载。但是，由于该 API 级别较低（因此用例范围更广泛），所以不会返回 Keras 模型。相反，它返回一个对象，其中包含可用于进行推断的函数。例如：<jupyter_code>DEFAULT_FUNCTION_KEY = "serving_default"
loaded = tf.saved_model.load(saved_model_path)
inference_func = loaded.signatures[DEFAULT_FUNCTION_KEY]<jupyter_output><empty_output><jupyter_text>加载的对象可能包含多个函数，每个函数与一个键关联。`"serving_default"` 是使用已保存的 Keras 模型的推断函数的默认键。要使用此函数进行推断，请运行以下代码：<jupyter_code>predict_dataset = eval_dataset.map(lambda image, label: image)
for batch in predict_dataset.take(1):
  print(inference_func(batch))<jupyter_output><empty_output><jupyter_text>您还可以采用分布式方式加载和进行推断：<jupyter_code>another_strategy = tf.distribute.MirroredStrategy()
with another_strategy.scope():
  loaded = tf.saved_model.load(saved_model_path)
  inference_func = loaded.signatures[DEFAULT_FUNCTION_KEY]

  dist_predict_dataset = another_strategy.experimental_distribute_dataset(
      predict_dataset)

  # Calling the function in a distributed manner
  for batch in dist_predict_dataset:
    another_strategy.run(inference_func,args=(batch,))<jupyter_output><empty_output><jupyter_text>调用已恢复的函数只是基于已保存模型的前向传递（预测）。如果您想继续训练加载的函数，或者将加载的函数嵌入到更大的模型中，应如何操作？ 通常的做法是将此加载对象包装到 Keras 层以实现此目的。幸运的是，[TF Hub](https://tensorflow.google.cn/hub) 为此提供了 [hub.KerasLayer](https://github.com/tensorflow/hub/blob/master/tensorflow_hub/keras_layer.py)，如下所示：<jupyter_code>import tensorflow_hub as hub

def build_model(loaded):
  x = tf.keras.layers.Input(shape=(28, 28, 1), name='input_x')
  # Wrap what's loaded to a KerasLayer
  keras_layer = hub.KerasLayer(loaded, trainable=True)(x)
  model = tf.keras.Model(x, keras_layer)
  return model

another_strategy = tf.distribute.MirroredStrategy()
with another_strategy.scope():
  loaded = tf.saved_model.load(saved_model_path)
  model = build_model(loaded)

  model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                optimizer=tf.keras.optimizers.Adam(),
                metrics=['accuracy'])
  model.fit(train_dataset, epochs=2)<jupyter_output><empty_output><jupyter_text>如您所见，`hub.KerasLayer` 可将从 `tf.saved_model.load()` 加载回的结果包装到可供构建其他模型的 Keras 层。这对于迁移学习非常实用。   我应使用哪种 API？ 对于保存，如果您使用的是 Keras 模型，那么始终建议使用 Keras 的 `model.save()` API。如果您所保存的不是 Keras 模型，那么您只能选择使用较低级的 API。对于加载，使用哪种 API 取决于您要从加载的 API 中获得什么。如果您无法或不想获取 Keras 模型，请使用 `tf.saved_model.load()`。否则，请使用 `tf.keras.models.load_model()`。请注意，只有保存 Keras 模型后，才能恢复 Keras 模型。可以混合使用 API。您可以使用 `model.save` 保存 Keras 模型，并使用低级 API `tf.saved_model.load` 加载非 Keras 模型。<jupyter_code>model = get_model()

# Saving the model using Keras's save() API
model.save(keras_model_path) 

another_strategy = tf.distribute.MirroredStrategy()
# Loading the model using lower level API
with another_strategy.scope():
  loaded = tf.saved_model.load(keras_model_path)<jupyter_output><empty_output><jupyter_text>警告 有一种特殊情况，您的 Keras 模型没有明确定义的输入。例如，可以创建没有任何输入形状的序贯模型 (`Sequential([Dense(3), ...]`)。子类化模型在初始化后也没有明确定义的输入。在这种情况下，在保存和加载时都应坚持使用较低级别的 API，否则会出现错误。要检查您的模型是否具有明确定义的输入，只需检查 `model.inputs` 是否为 `None`。如果非 `None`，则一切正常。在 `.fit`、`.evaluate`、`.predict` 中使用模型，或调用模型 (`model(inputs)`) 时，输入形状将自动定义。以下为示例：<jupyter_code>class SubclassedModel(tf.keras.Model):

  output_name = 'output_layer'

  def __init__(self):
    super(SubclassedModel, self).__init__()
    self._dense_layer = tf.keras.layers.Dense(
        5, dtype=tf.dtypes.float32, name=self.output_name)

  def call(self, inputs):
    return self._dense_layer(inputs)

my_model = SubclassedModel()
# my_model.save(keras_model_path)  # ERROR! 
tf.saved_model.save(my_model, saved_model_path)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11669539848826416, 'fraction_numerical': 0.006497811961278345, 'mean_word_length': 6.7592592592592595, 'pattern_counts': {'":': 0, '<': 57, '<?xml version=': 0, '>': 57, 'https://': 3, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '14', 'dir': 'jupyter-structured-clean-dedup', 'id': '9377', 'license': 'Apache-2.0', 'n_tokens_mistral': 3390, 'n_tokens_neox': 3483, 'n_words': 438, 'path': 'site/zh-cn/tutorials/distribute/save_and_load.ipynb', 'repo_name': 'FontTian/docs-l10n', 'size': '7541'}